Add max_concurrency feature to ExtraTaskGroup

Author: jsbueno

README.md

@@ -44,8 +44,17 @@ any exception is raised in any of the taskgroup tasks,
 all sibling incomplete tasks get cancelled immediatelly.
 
 With ExtraTaskGroup, all created tasks are run to completion,
-and any exceptions are bubbled up as ExceptionGroups on
-the host task.
+by default and any exceptions are bubbled up as ExceptionGroups on
+the host task. If the classic "cancel all other tasks" behavior
+is desired, the named argument `default_abort=True` can be
+passed when the group is created.
+
+ExtraTaskGroup will also optionally limit the number
+of concurrent tasks that are executed. If `max_concurrency` is given,
+tasks created with an inner `.create_task` call will be
+bounded by a Semaphore, meaning at most "max_concurrency" tasks
+will be running at the same time, and others will start running as the first
+ones are completed.
 
 ```python
 import asyncio
@@ -70,6 +79,9 @@ asyncio.run(main())
 
 ```
 
+new in 0.4: the max_concurrency and default_abort paramters
+
+
 sync_to_async
 ----------------------
 Allows calling an async function from a synchronous context.

extraasync/__init__.py

@@ -3,7 +3,7 @@
 from .sync_async_bridge import sync_to_async, async_to_sync
 from .async_hooks import at_loop_stop_callback, remove_loop_stop_callback
 
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 
 
 __all__ = ["aenumerate", "ExtraTaskGroup", "sync_to_async", "async_to_sync", "at_loop_stop_callback", "remove_loop_stop_callback",  "__version__"]

extraasync/taskgroup.py

@@ -1,5 +1,9 @@
+import asyncio
 from asyncio import TaskGroup
 
+from typing import Optional
+
+
 # Idea originally developed for an answer on StackOverflow
 # at: https://stackoverflow.com/questions/75250788/how-to-prevent-python3-11-taskgroup-from-canceling-all-the-tasks/75261668#75261668
 
@@ -17,7 +21,7 @@
 
 
 class ExtraTaskGroup(TaskGroup):
-    def __init__(self, *, default_abort: bool = False):
+    def __init__(self, *, max_concurrency: Optional[int] = None, default_abort: bool = False):
         """A subclass of asyncio.TaskGroup
 
         By default, the different behavior is that if a
@@ -33,16 +37,34 @@ def __init__(self, *, default_abort: bool = False):
         single exception.
 
         Args:
+            max_concurrency: If given, the maximum number of active spawned tasks permited at each time.
+                The asynchronous "acreate_task" method should be used then, and calling "create_task" will raise
+                an error.
+                Defaults to None. implemented internally as using a semaphore.
+
             default_abort: if True, allows the default asyncio.TaskGroup behavior
-            or aborting all other running tasks when the first one raises an exception.
+                or aborting all other running tasks when the first one raises an exception.
 
 
         """
-        self.default_abort = default_abort
+        self.__max_concurrency = max_concurrency
+        if max_concurrency != None:
+            self.__semaphore = asyncio.Semaphore(max_concurrency)
+        self.__default_abort = default_abort
         super().__init__()
 
+
+    def create_task(self, coro, *args, **kwargs):
+        if self.__max_concurrency is None:
+            return super().create_task(coro, *args, **kwargs)
+        return super().create_task(self.__managed_task(coro), *args, **kwargs)
+
+    async def __managed_task(self, coro):
+        async with self.__semaphore:
+            return await coro
+
     def _abort(self):
-        if self.default_abort:
+        if self.__default_abort:
             return super()._abort()
         return None
 

tests/test_taskgroup.py

@@ -66,3 +66,30 @@ class Dummy:
 
 
 
+@pytest.mark.parametrize(
+    ["max_tasks", "max_concurrency"],[
+        (2, 1),
+        (5, 1),
+        (5, 2),
+        (30, 10),
+])
+@pytest.mark.asyncio
+async def test_extrataskgroup_limits_concurrency(max_tasks, max_concurrency):
+    running = 0
+    max_running = 0
+    async def blah():
+        nonlocal max_running, running
+        running += 1
+        max_running = max(max_running, running)
+        await asyncio.sleep(0.05)
+        running -= 1
+        return 23
+
+    tasks = set()
+
+    async with ExtraTaskGroup(max_concurrency=max_concurrency) as tg:
+        for _ in range(max_tasks):
+            tasks.add(tg.create_task(blah()))
+
+    assert all(task.result() == 23 for task in tasks)
+    assert max_running == max_concurrency