clean frontend disconnection protocol

yacv_server/pubsub.py

@@ -1,4 +1,4 @@
-import threading
+import queue
 import queue
 import threading
 from typing import List, TypeVar, \
@@ -8,6 +8,8 @@ from yacv_server.mylogger import logger
 
 T = TypeVar('T')
 
+_end_of_queue = object()
+
 
 class BufferedPubSub(Generic[T]):
     """A simple implementation of publish-subscribe pattern using threading and buffering all previous events"""
@@ -45,7 +47,7 @@ class BufferedPubSub(Generic[T]):
                 for event in self._buffer:
                     q.put(event)
             if not include_future:
-                q.put(None)
+                q.put(_end_of_queue)
         return q
 
     def _unsubscribe(self, q: queue.Queue[T]):
@@ -54,14 +56,18 @@ class BufferedPubSub(Generic[T]):
             self._subscribers.remove(q)
         logger.debug(f"Unsubscribed from %s (%d subscribers)", self, len(self._subscribers))
 
-    def subscribe(self, include_buffered: bool = True, include_future: bool = True) -> Generator[T, None, None]:
+    def subscribe(self, include_buffered: bool = True, include_future: bool = True, yield_timeout: float = 0.0) -> \
+            Generator[T, None, None]:
         """Subscribes to events as an generator that yields events and automatically unsubscribes"""
         q = self._subscribe(include_buffered, include_future)
         try:
             while True:
-                v = q.get()
+                try:
+                    v = q.get(timeout=yield_timeout)
+                except queue.Empty:
+                    v = None
                 # include_future is incompatible with None values as they are used to signal the end of the stream
-                if v is None and not include_future:
+                if v is _end_of_queue:
                     break
                 yield v
         finally:  # When aclose() is called
@@ -80,4 +86,4 @@ class BufferedPubSub(Generic[T]):
     def clear(self):
         """Clears the buffer"""
         with self._buffer_lock:
-            self._buffer.clear()
+            self._buffer.clear()