setup.cfg

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

src/asyncio_taskpool/pool.py

@@ -3,7 +3,6 @@ from asyncio import gather
 from asyncio.coroutines import iscoroutine, iscoroutinefunction
 from asyncio.exceptions import CancelledError
 from asyncio.locks import Event, Semaphore
-from asyncio.queues import Queue, QueueEmpty
 from asyncio.tasks import Task, create_task
 from functools import partial
 from math import inf
@@ -40,7 +39,8 @@ class BaseTaskPool:
         self._num_ended: int = 0
         self._idx: int = self._add_pool(self)
         self._name: str = name
-        self._before_gathering: List[Awaitable] = []
+        self._all_tasks_known_flag: Event = Event()
+        self._all_tasks_known_flag.set()
         self._interrupt_flag: Event = Event()
         log.debug("%s initialized", str(self))
 
@@ -201,13 +201,12 @@ class BaseTaskPool:
                 If `True`, even if the pool is closed, the task will still be started.
             end_callback (optional):
                 A callback to execute after the task has ended.
-                It is run with the task's ID as its only positional argument.
+                It is run with the `task_id` as its only positional argument.
             cancel_callback (optional):
                 A callback to execute after cancellation of the task.
-                It is run with the task's ID as its only positional argument.
+                It is run with the `task_id` as its only positional argument.
 
         Raises:
-            `asyncio_taskpool.exceptions.NotCoroutine` if `awaitable` is not a coroutine.
             `asyncio_taskpool.exceptions.PoolIsClosed` if the pool has been closed and `ignore_closed` is `False`.
         """
         if not iscoroutine(awaitable):
@@ -330,7 +329,7 @@ class BaseTaskPool:
         """
         if self._open:
             raise exceptions.PoolStillOpen("Pool must be closed, before tasks can be gathered")
-        await gather(*self._before_gathering)
+        await self._all_tasks_known_flag.wait()
         results = await gather(*self._ended.values(), *self._cancelled.values(), *self._running.values(),
                                return_exceptions=return_exceptions)
         self._ended = self._cancelled = self._running = {}
@@ -340,270 +339,68 @@ class BaseTaskPool:
 
 
 class TaskPool(BaseTaskPool):
-    """
-    General task pool class.
-    Attempts to somewhat emulate part of the interface of `multiprocessing.pool.Pool` from the stdlib.
-
-    A `TaskPool` instance can manage an arbitrary number of concurrent tasks from any coroutine function.
-    Tasks in the pool can all belong to the same coroutine function,
-    but they can also come from any number of different and unrelated coroutine functions.
-
-    As long as there is room in the pool, more tasks can be added. (By default, there is no pool size limit.)
-    Each task started in the pool receives a unique ID, which can be used to cancel specific tasks at any moment.
-
-    Adding tasks blocks **only if** the pool is full at that moment.
-    """
-
     async def _apply_one(self, func: CoroutineFunc, args: ArgsT = (), kwargs: KwArgsT = None,
                          end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> int:
-        """
-        Creates a coroutine with the supplied arguments and runs it as a new task in the pool.
-
-        This method blocks, **only if** the pool is full.
-
-        Args:
-            func:
-                The coroutine function to be run as a task within the task pool.
-            args (optional):
-                The positional arguments to pass into the function call.
-            kwargs (optional):
-                The keyword-arguments to pass into the function call.
-            end_callback (optional):
-                A callback to execute after the task has ended.
-                It is run with the task's ID as its only positional argument.
-            cancel_callback (optional):
-                A callback to execute after cancellation of the task.
-                It is run with the task's ID as its only positional argument.
-
-        Returns:
-            The newly spawned task's ID within the pool.
-        """
         if kwargs is None:
             kwargs = {}
         return await self._start_task(func(*args, **kwargs), end_callback=end_callback, cancel_callback=cancel_callback)
 
     async def apply(self, func: CoroutineFunc, args: ArgsT = (), kwargs: KwArgsT = None, num: int = 1,
                     end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> List[int]:
-        """
-        Creates an arbitrary number of coroutines with the supplied arguments and runs them as new tasks in the pool.
-        Each coroutine looks like `func(*args, **kwargs)`.
-
-        This method blocks, **only if** there is not enough room in the pool for the desired number of new tasks.
-
-        Args:
-            func:
-                The coroutine function to use for spawning the new tasks within the task pool.
-            args (optional):
-                The positional arguments to pass into each function call.
-            kwargs (optional):
-                The keyword-arguments to pass into each function call.
-            num (optional):
-                The number of tasks to spawn with the specified parameters.
-            end_callback (optional):
-                A callback to execute after a task has ended.
-                It is run with the task's ID as its only positional argument.
-            cancel_callback (optional):
-                A callback to execute after cancellation of a task.
-                It is run with the task's ID as its only positional argument.
-
-        Returns:
-            The newly spawned tasks' IDs within the pool as a list of integers.
-
-        Raises:
-            `NotCoroutine` if `func` is not a coroutine function.
-            `PoolIsClosed` if the pool has been closed already.
-        """
         ids = await gather(*(self._apply_one(func, args, kwargs, end_callback, cancel_callback) for _ in range(num)))
         # TODO: for some reason PyCharm wrongly claims that `gather` returns a tuple of exceptions
         assert isinstance(ids, list)
         return ids
 
-    async def _queue_producer(self, q: Queue, args_iter: Iterator[Any]) -> None:
+    async def _next_callback(self, task_id: int, func: CoroutineFunc, args_iter: Iterator[Any], arg_stars: int = 0,
-        """
-        Keeps the arguments queue from `_map()` full as long as the iterator has elements.
-        If the `_interrupt_flag` gets set, the loop ends prematurely.
-
-        Args:
-            q:
-                The queue of function arguments to consume for starting the next task.
-            args_iter:
-                The iterator of function arguments to put into the queue.
-        """
-        for arg in args_iter:
-            if self._interrupt_flag.is_set():
-                break
-            await q.put(arg)  # This blocks as long as the queue is full.
-
-    async def _queue_consumer(self, q: Queue, func: CoroutineFunc, arg_stars: int = 0,
                              end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
-        """
+        reached_end = await self._start_next_task(func, args_iter, arg_stars=arg_stars,
-        Wrapper around the `_start_task()` taking the next element from the arguments queue set up in `_map()`.
+                                                  end_callback=end_callback, cancel_callback=cancel_callback)
-        Partially constructs the `_queue_callback` function with the same arguments.
+        if reached_end:
+            self._all_tasks_known_flag.set()
+        await execute_optional(end_callback, args=(task_id,))
 
-        Args:
+    async def _start_next_task(self, func: CoroutineFunc, args_iter: Iterator[Any], arg_stars: int = 0,
-            q:
+                               end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> bool:
-                The queue of function arguments to consume for starting the next task.
+        if self._interrupt_flag.is_set():
-            func:
+            return True
-                The coroutine function to use for spawning the tasks within the task pool.
-            arg_stars (optional):
-                Whether or not to unpack an element from `q` using stars; must be 0, 1, or 2.
-            end_callback (optional):
-                The actual callback specified to execute after the task (and the next one) has ended.
-                It is run with the task's ID as its only positional argument.
-            cancel_callback (optional):
-                The callback that was specified to execute after cancellation of the task (and the next one).
-                It is run with the task's ID as its only positional argument.
-        """
-        try:
-            arg = q.get_nowait()
-        except QueueEmpty:
-            return
         try:
             await self._start_task(
-                star_function(func, arg, arg_stars=arg_stars),
+                star_function(func, next(args_iter), arg_stars=arg_stars),
                 ignore_closed=True,
-                end_callback=partial(TaskPool._queue_callback, self, q=q, func=func, arg_stars=arg_stars,
+                end_callback=partial(TaskPool._next_callback, self, func=func, args_iter=args_iter, arg_stars=arg_stars,
                                      end_callback=end_callback, cancel_callback=cancel_callback),
                 cancel_callback=cancel_callback
             )
-        finally:
+        except StopIteration:
-            q.task_done()
+            return True
-
+        return False
-    async def _queue_callback(self, task_id: int, q: Queue, func: CoroutineFunc, arg_stars: int = 0,
-                              end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
-        """
-        Wrapper around an end callback function passed into the `_map()` method.
-        Triggers the next `_queue_consumer` with the same arguments.
-
-        Args:
-            task_id:
-                The ID of the ending task.
-            q:
-                The queue of function arguments to consume for starting the next task.
-            func:
-                The coroutine function to use for spawning the tasks within the task pool.
-            arg_stars (optional):
-                Whether or not to unpack an element from `q` using stars; must be 0, 1, or 2.
-            end_callback (optional):
-                The actual callback specified to execute after the task (and the next one) has ended.
-                It is run with the `task_id` as its only positional argument.
-            cancel_callback (optional):
-                The callback that was specified to execute after cancellation of the task (and the next one).
-                It is run with the `task_id` as its only positional argument.
-        """
-        await self._queue_consumer(q, func, arg_stars, end_callback=end_callback, cancel_callback=cancel_callback)
-        await execute_optional(end_callback, args=(task_id,))
 
     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.
-        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.
-
-        It sets up an internal queue which is filled while consuming the arguments iterable.
-        The queue's `join()` method is added to the pool's `_before_gathering` list.
-
-        Args:
-            func:
-                The coroutine function to use for spawning the new tasks within the task pool.
-            args_iter:
-                The iterable of arguments; each element is to be passed into a `func` call when spawning a new task.
-            arg_stars (optional):
-                Whether or not to unpack an element from `args_iter` using stars; must be 0, 1, or 2.
-            num_tasks (optional):
-                The maximum number of the new tasks to run concurrently.
-            end_callback (optional):
-                A callback to execute after a task has ended.
-                It is run with the task's ID as its only positional argument.
-            cancel_callback (optional):
-                A callback to execute after cancellation of a task.
-                It is run with the task's ID as its only positional argument.
-
-        Raises:
-            `asyncio_taskpool.exceptions.PoolIsClosed` if the pool has been closed.
-        """
         if not self.is_open:
             raise exceptions.PoolIsClosed("Cannot start new tasks")
-        args_queue = Queue(maxsize=num_tasks)
+        if self._all_tasks_known_flag.is_set():
-        self._before_gathering.append(args_queue.join())
+            self._all_tasks_known_flag.clear()
         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))
         for _ in range(num_tasks):
-            # This is where blocking can occur, if the pool is full.
+            reached_end = await self._start_next_task(func, args_iter, arg_stars, end_callback, cancel_callback)
-            await self._queue_consumer(args_queue, func,
+            if reached_end:
-                                       arg_stars=arg_stars, end_callback=end_callback, cancel_callback=cancel_callback)
+                self._all_tasks_known_flag.set()
+                break
 
-    async def map(self, func: CoroutineFunc, arg_iter: ArgsT, num_tasks: int = 1,
+    async def map(self, func: CoroutineFunc, args_iter: ArgsT, num_tasks: int = 1,
                   end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
-        """
+        await self._map(func, args_iter, arg_stars=0, num_tasks=num_tasks,
-        An asyncio-task-based equivalent of the `multiprocessing.pool.Pool.map` method.
-
-        Creates coroutines with arguments from a supplied iterable and runs them as new tasks in the pool in batches.
-        Each coroutine looks like `func(arg)`, `arg` being an element from the iterable.
-
-        Once the first batch of tasks has started to run, this method returns.
-        As soon as on of them finishes, it triggers the start of a new task (assuming there is room in the pool)
-        consuming the next element from the arguments iterable.
-        If the size of the pool never imposes a limit, this ensures that there is almost continuously the desired number
-        of tasks from this call concurrently running within the pool.
-
-        This method blocks, **only if** there is not enough room in the pool for the first batch of new tasks.
-
-        Args:
-            func:
-                The coroutine function to use for spawning the new tasks within the task pool.
-            arg_iter:
-                The iterable of arguments; each argument is to be passed into a `func` call when spawning a new task.
-            num_tasks (optional):
-                The maximum number of the new tasks to run concurrently.
-            end_callback (optional):
-                A callback to execute after a task has ended.
-                It is run with the task's ID as its only positional argument.
-            cancel_callback (optional):
-                A callback to execute after cancellation of a task.
-                It is run with the task's ID as its only positional argument.
-
-        Raises:
-            `PoolIsClosed` if the pool has been closed.
-            `NotCoroutine` if `func` is not a coroutine function.
-        """
-        await self._map(func, arg_iter, arg_stars=0, num_tasks=num_tasks,
                         end_callback=end_callback, cancel_callback=cancel_callback)
 
     async def starmap(self, func: CoroutineFunc, args_iter: Iterable[ArgsT], num_tasks: int = 1,
                       end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
-        """
-        Like `map()` except that the elements of `args_iter` are expected to be iterables themselves to be unpacked as
-        positional arguments to the function.
-        Each coroutine then looks like `func(*arg)`, `arg` being an element from `args_iter`.
-        """
         await self._map(func, args_iter, arg_stars=1, num_tasks=num_tasks,
                         end_callback=end_callback, cancel_callback=cancel_callback)
 
     async def doublestarmap(self, func: CoroutineFunc, kwargs_iter: Iterable[KwArgsT], num_tasks: int = 1,
                             end_callback: EndCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
-        """
-        Like `map()` except that the elements of `kwargs_iter` are expected to be iterables themselves to be unpacked as
-        keyword-arguments to the function.
-        Each coroutine then looks like `func(**arg)`, `arg` being an element from `kwargs_iter`.
-        """
         await self._map(func, kwargs_iter, arg_stars=2, num_tasks=num_tasks,
                         end_callback=end_callback, cancel_callback=cancel_callback)