factored out the queue setup of TaskPool._map

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = asyncio-taskpool
-version = 0.1.1
+version = 0.1.3
 author = Daniil Fajnberg
 author_email = mail@daniil.fajnberg.de
 description = Dynamically manage pools of asyncio tasks

src/asyncio_taskpool/helpers.py

@@ -1,4 +1,5 @@
 from asyncio.coroutines import iscoroutinefunction
+from asyncio.queues import Queue
 from typing import Any, Optional
 
 from .types import T, AnyCallableT, ArgsT, KwArgsT
@@ -22,3 +23,7 @@ def star_function(function: AnyCallableT, arg: Any, arg_stars: int = 0) -> T:
     if arg_stars == 2:
         return function(**arg)
     raise ValueError(f"Invalid argument arg_stars={arg_stars}; must be 0, 1, or 2.")
+
+
+async def join_queue(q: Queue) -> None:
+    await q.join()

src/asyncio_taskpool/pool.py

@@ -10,7 +10,7 @@ from math import inf
 from typing import Any, Awaitable, Dict, Iterable, Iterator, List
 
 from . import exceptions
-from .helpers import execute_optional, star_function
+from .helpers import execute_optional, star_function, join_queue
 from .types import ArgsT, KwArgsT, CoroutineFunc, EndCallbackT, CancelCallbackT
 
 
@@ -297,7 +297,8 @@ class BaseTaskPool:
             return_exceptions (optional): Passed directly into `gather`.
         """
         results = await gather(*self._ended.values(), *self._cancelled.values(), return_exceptions=return_exceptions)
-        self._ended = self._cancelled = {}
+        self._ended.clear()
+        self._cancelled.clear()
         if self._interrupt_flag.is_set():
             self._interrupt_flag.clear()
         return results
@@ -333,7 +334,10 @@ class BaseTaskPool:
         await gather(*self._before_gathering)
         results = await gather(*self._ended.values(), *self._cancelled.values(), *self._running.values(),
                                return_exceptions=return_exceptions)
-        self._ended = self._cancelled = self._running = {}
+        self._ended.clear()
+        self._cancelled.clear()
+        self._running.clear()
+        self._before_gathering.clear()
         if self._interrupt_flag.is_set():
             self._interrupt_flag.clear()
         return results
@@ -494,10 +498,32 @@ class TaskPool(BaseTaskPool):
         await self._queue_consumer(q, func, arg_stars, end_callback=end_callback, cancel_callback=cancel_callback)
         await execute_optional(end_callback, args=(task_id,))
 
+    def _fill_args_queue(self, q: Queue, args_iter: ArgsT, num_tasks: int) -> int:
+        args_iter = iter(args_iter)
+        try:
+            # Here we guarantee that the queue will contain as many arguments as needed for starting the first batch of
+            # tasks, which will be at most `num_tasks` (meaning the queue will be full).
+            for i in range(num_tasks):
+                q.put_nowait(next(args_iter))
+        except StopIteration:
+            # If we get here, this means that the number of elements in the arguments iterator was less than the
+            # specified `num_tasks`. Thus, the number of tasks to start immediately will be the size of the queue.
+            # The `_queue_producer` won't be necessary, since we already put all the elements in the queue.
+            num_tasks = q.qsize()
+        else:
+            # There may be more elements in the arguments iterator, so we need the `_queue_producer`.
+            # It will have exclusive access to the `args_iter` from now on.
+            # If the queue is full already, it will wait until one of the tasks in the first batch ends, before putting
+            # the next item in it.
+            create_task(self._queue_producer(q, args_iter))
+        return num_tasks
+
     async def _map(self, func: CoroutineFunc, args_iter: ArgsT, arg_stars: int = 0, num_tasks: int = 1,
                    end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
         """
         Creates coroutines with arguments from a supplied iterable and runs them as new tasks in the pool in batches.
+        TODO: If task groups are implemented, consider adding all tasks from one call of this method to the same group
+              and referring to "group size" rather than chunk/batch size.
         Each coroutine looks like `func(arg)`, `func(*arg)`, or `func(**arg)`, `arg` being an element from the iterable.
 
         This method blocks, **only if** there is not enough room in the pool for the first batch of new tasks.
@@ -527,24 +553,8 @@ class TaskPool(BaseTaskPool):
         if not self.is_open:
             raise exceptions.PoolIsClosed("Cannot start new tasks")
         args_queue = Queue(maxsize=num_tasks)
-        self._before_gathering.append(args_queue.join())
-        args_iter = iter(args_iter)
-        try:
-            # Here we guarantee that the queue will contain as many arguments as needed for starting the first batch of
-            # tasks, which will be at most `num_tasks` (meaning the queue will be full).
-            for i in range(num_tasks):
-                args_queue.put_nowait(next(args_iter))
-        except StopIteration:
-            # If we get here, this means that the number of elements in the arguments iterator was less than the
-            # specified `num_tasks`. Thus, the number of tasks to start immediately will be the size of the queue.
-            # The `_queue_producer` won't be necessary, since we already put all the elements in the queue.
-            num_tasks = args_queue.qsize()
-        else:
-            # There may be more elements in the arguments iterator, so we need the `_queue_producer`.
-            # It will have exclusive access to the `args_iter` from now on.
-            # If the queue is full already, it will wait until one of the tasks in the first batch ends, before putting
-            # the next item in it.
-            create_task(self._queue_producer(args_queue, args_iter))
+        self._before_gathering.append(join_queue(args_queue))
+        num_tasks = self._fill_args_queue(args_queue, args_iter, num_tasks)
         for _ in range(num_tasks):
             # This is where blocking can occur, if the pool is full.
             await self._queue_consumer(args_queue, func,

tests/test_pool.py

@@ -67,8 +67,7 @@ class BaseTaskPoolTestCase(IsolatedAsyncioTestCase):
         self.assertEqual(0, self.task_pool._num_ended)
         self.assertEqual(self.mock_idx, self.task_pool._idx)
         self.assertEqual(self.test_pool_name, self.task_pool._name)
-        self.assertIsInstance(self.task_pool._all_tasks_known_flag, asyncio.locks.Event)
-        self.assertTrue(self.task_pool._all_tasks_known_flag.is_set())
+        self.assertListEqual(self.task_pool._before_gathering, EMPTY_LIST)
         self.assertIsInstance(self.task_pool._interrupt_flag, asyncio.locks.Event)
         self.assertFalse(self.task_pool._interrupt_flag.is_set())
         self.mock__add_pool.assert_called_once_with(self.task_pool)
@@ -342,11 +341,11 @@ class BaseTaskPoolTestCase(IsolatedAsyncioTestCase):
         self.assertFalse(self.task_pool._open)
 
     async def test_gather(self):
-        mock_wait = AsyncMock()
-        self.task_pool._all_tasks_known_flag = MagicMock(wait=mock_wait)
         test_exception = TestException()
         mock_ended_func, mock_cancelled_func = AsyncMock(return_value=FOO), AsyncMock(side_effect=test_exception)
         mock_running_func = AsyncMock(return_value=BAR)
+        mock_queue_join = AsyncMock()
+        self.task_pool._before_gathering = before_gather = [mock_queue_join()]
         self.task_pool._ended = ended = {123: mock_ended_func()}
         self.task_pool._cancelled = cancelled = {456: mock_cancelled_func()}
         self.task_pool._running = running = {789: mock_running_func()}
@@ -358,25 +357,23 @@ class BaseTaskPoolTestCase(IsolatedAsyncioTestCase):
         self.assertDictEqual(self.task_pool._ended, ended)
         self.assertDictEqual(self.task_pool._cancelled, cancelled)
         self.assertDictEqual(self.task_pool._running, running)
+        self.assertListEqual(self.task_pool._before_gathering, before_gather)
         self.assertTrue(self.task_pool._interrupt_flag.is_set())
-        mock_wait.assert_not_awaited()
 
         self.task_pool._open = False
 
-        def check_assertions() -> None:
+        def check_assertions(output) -> None:
             self.assertListEqual([FOO, test_exception, BAR], output)
             self.assertDictEqual(self.task_pool._ended, EMPTY_DICT)
             self.assertDictEqual(self.task_pool._cancelled, EMPTY_DICT)
             self.assertDictEqual(self.task_pool._running, EMPTY_DICT)
+            self.assertListEqual(self.task_pool._before_gathering, EMPTY_LIST)
             self.assertFalse(self.task_pool._interrupt_flag.is_set())
-            mock_wait.assert_awaited_once_with()
 
-        output = await self.task_pool.gather(return_exceptions=True)
-        check_assertions()
-        mock_wait.reset_mock()
+        check_assertions(await self.task_pool.gather(return_exceptions=True))
 
+        self.task_pool._before_gathering = [mock_queue_join()]
         self.task_pool._ended = {123: mock_ended_func()}
         self.task_pool._cancelled = {456: mock_cancelled_func()}
         self.task_pool._running = {789: mock_running_func()}
-        output = await self.task_pool.gather(return_exceptions=True)
-        check_assertions()
+        check_assertions(await self.task_pool.gather(return_exceptions=True))