queue: added task_done/join behavior and test (asynctest)

Author: stephanelsmith

v3/primitives/queue.py

@@ -26,6 +26,9 @@ 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()
+
     def _get(self):
         self._evget.set()  # Schedule all tasks waiting on get
         self._evget.clear()
@@ -45,6 +48,7 @@ def get_nowait(self):  # Remove and return an item from the queue.
         return self._get()
 
     def _put(self, val):
+        self._jncnt += 1
         self._evput.set()  # Schedule tasks waiting on put
         self._evput.clear()
         self._queue.append(val)
@@ -71,3 +75,15 @@ 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 task_done(self):
+        self._jncnt -= 1
+        if self._jncnt <= 0:
+            self._jnevt.set()
+        else:
+            self._jnevt.clear()
+
+    async def join(self):
+        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,68 @@ 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_producer(q):
+    print('producer','loading jobs')
+    for x in range(10):
+        await q.put(x)
+    print('producer','await q.join')
+    await q.join()
+    print('producer','joined!', 'task done!')
+async def q_task_done_join_go():
+    q = Queue()
+
+    consumer_task = asyncio.create_task(q_task_done_join_consumer(q))
+    producer_task = asyncio.create_task(q_task_done_join_producer(q))
+    await asyncio.sleep(0)
+
+    print('test','await q.join')
+    await q.join()
+    print('test','all jobs done!')
+
+    print('test','join again')
+    await q.join()
+
+    await asyncio.sleep(0)
+    print('test','producer_task.done()?', producer_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
+producer loading jobs
+producer await q.join
+test await q.join
+consumer got/processing 0
+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
+producer joined! task done!
+test all jobs done!
+test join again
+test producer_task.done()? True
+test DONE
+''', 5)
+    asyncio.run(q_task_done_join_go())
+
+
 # ************ ************
 def test(n):
     try:
@@ -657,6 +720,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: