Merge pull request peterhinch#109 from stephanelsmith/master

Queue: Add task_done/join

Author: peterhinch

v3/docs/TUTORIAL.md

@@ -973,12 +973,15 @@ Synchronous methods (immediate return):
  queue is full.
  * `get_nowait` No arg. Returns an object from the queue. Raises an exception
  if the queue is empty.
+ * `task_done` No arg. Indicate that a task associated with a dequeued item is complete.
 
 Asynchronous methods:  
  * `put` Arg: the object to put on the queue. If the queue is full, it will
  block until space is available.
  * `get` No arg. Returns an object from the queue. If the queue is empty, it
  will block until an object is put on the queue.
+ * `join` No arg. Block until all items in the queue have been received and
+ processed (indicated via task_done).
 
 ```python
 import uasyncio as asyncio

v3/primitives/queue.py

@@ -26,6 +26,10 @@ def __init__(self, maxsize=0):
         self._evput = asyncio.Event()  # Triggered by put, tested by get
         self._evget = asyncio.Event()  # Triggered by get, tested by put
 
+        self._jncnt = 0
+        self._jnevt = asyncio.Event()
+        self._upd_jnevt(0) #update join event
+
     def _get(self):
         self._evget.set()  # Schedule all tasks waiting on get
         self._evget.clear()
@@ -45,6 +49,7 @@ def get_nowait(self):  # Remove and return an item from the queue.
         return self._get()
 
     def _put(self, val):
+        self._upd_jnevt(1) # update join event
         self._evput.set()  # Schedule tasks waiting on put
         self._evput.clear()
         self._queue.append(val)
@@ -71,3 +76,19 @@ def full(self):  # Return True if there are maxsize items in the queue.
         # Note: if the Queue was initialized with maxsize=0 (the default) or
         # any negative number, then full() is never True.
         return self.maxsize > 0 and self.qsize() >= self.maxsize
+
+
+    def _upd_jnevt(self, inc:int): # #Update join count and join event
+        self._jncnt += inc
+        if self._jncnt <= 0:
+            self._jnevt.set()
+        else:
+            self._jnevt.clear()
+
+    def task_done(self): # Task Done decrements counter
+        self._upd_jnevt(-1)
+
+    async def join(self): # Wait for join event
+        await self._jnevt.wait()
+
+

v3/primitives/tests/asyntest.py

@@ -34,6 +34,7 @@ def print_tests():
 test(7)  Test the Condition class.
 test(8)  Test the Queue class.
 test(9)  Test the RingbufQueue class.
+test(10) Test the Queue task_done/join behavior.
 '''
     print('\x1b[32m')
     print(st)
@@ -636,6 +637,70 @@ def rbq_test():
 ''', 6)
     asyncio.run(rbq_go())
 
+
+# ************ Queue task_done/join test ************
+async def q_task_done_join_consumer(q):
+    while True:
+        r = await q.get()
+        print('consumer', 'got/processing {}'.format(r))
+        await asyncio.sleep(.5)
+        q.task_done()
+async def q_task_done_join_waiter(q):
+    print('waiter','await q.join')
+    await q.join()
+    print('waiter','joined!', 'task done!')
+async def q_task_done_join_go():
+    q = Queue()
+
+    #empty queue should not block join
+    print('test', 'await empty q.join')
+    await q.join()
+    print('test', 'pass')
+
+    consumer_task = asyncio.create_task(q_task_done_join_consumer(q))
+    waiter_task = asyncio.create_task(q_task_done_join_waiter(q))
+
+    #add jobs
+    for x in range(10):
+        await q.put(x)
+
+    print('test','await q.join')
+    await q.join()
+    print('test','all jobs done!')
+
+    await asyncio.sleep(0)
+    print('test','waiter_task.done()?', waiter_task.done())
+
+    consumer_task.cancel()
+    await asyncio.gather(consumer_task, return_exceptions=True)
+
+    print('test','DONE')
+
+
+def q_task_done_join_test():
+    printexp('''Test Queue task_done/join behaviors
+test await empty q.join
+test pass
+test await q.join
+consumer got/processing 0
+waiter await q.join
+consumer got/processing 1
+consumer got/processing 2
+consumer got/processing 3
+consumer got/processing 4
+consumer got/processing 5
+consumer got/processing 6
+consumer got/processing 7
+consumer got/processing 8
+consumer got/processing 9
+test all jobs done!
+waiter joined! task done!
+test waiter_task.done()? True
+test DONE
+''', 5)
+    asyncio.run(q_task_done_join_go())
+
+
 # ************ ************
 def test(n):
     try:
@@ -657,6 +722,8 @@ def test(n):
             queue_test()  # Test the Queue class.
         elif n == 9:
             rbq_test()  # Test the RingbufQueue class.
+        elif n == 10:
+            q_task_done_join_test()  # Test the Queue task_done/join behavior.
     except KeyboardInterrupt:
         print('Interrupted')
     finally: