Skip to content

Commit

Permalink
Store reference to loop in BGTasks
Browse files Browse the repository at this point in the history
resolves #9

Skips asyncio.get_running_loop overhead for task creation.
  • Loading branch information
mikeshardmind committed Feb 4, 2025
1 parent 0fcb39a commit 2c0b8ad
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 3 deletions.
2 changes: 1 addition & 1 deletion src/async_utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
__author__ = "Michael Hall"
__license__ = "Apache-2.0"
__copyright__ = "Copyright 2020-Present Michael Hall"
__version__ = "2025.01.30"
__version__ = "2025.02.04"

import os
import sys
Expand Down
26 changes: 24 additions & 2 deletions src/async_utils/bg_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,21 @@
from . import _typings as t

type _CoroutineLike[T] = Coroutine[t.Any, t.Any, T]
type _LoopLike = asyncio.AbstractEventLoop | _UnboundLoopSentinel


__all__ = ("BGTasks",)


class _UnboundLoopSentinel:
def create_task(*args: object, **kwargs: object) -> t.Never:
msg = """
BGTasks is intended for use as a context manager. Using create_task
prior to entering the context is not supported.
"""
raise RuntimeError(msg)


class BGTasks:
"""An intentionally dumber task group.
Expand All @@ -42,6 +52,7 @@ class BGTasks:
def __init__(self, exit_timeout: float | None) -> None:
self._tasks: set[asyncio.Task[object]] = set()
self._etime: float | None = exit_timeout
self._loop: _LoopLike = _UnboundLoopSentinel()

def create_task[T](
self,
Expand All @@ -50,18 +61,29 @@ def create_task[T](
name: str | None = None,
context: Context | None = None,
) -> asyncio.Task[T]:
"""Create a task bound managed by this context manager.
"""Create a task attached to this context manager.
Returns
-------
asyncio.Task: The task that was created.
"""
t = asyncio.create_task(coro, name=name, context=context)
t = self._loop.create_task(coro, name=name, context=context)
if name is not None:
# See: python/cpython#113050
# PYUPDATE: remove this block at python 3.13 minimum
try:
set_name = t.set_name
except AttributeError:
pass
else:
set_name(name)

self._tasks.add(t)
t.add_done_callback(self._tasks.discard)
return t

async def __aenter__(self: t.Self) -> t.Self:
self._loop = asyncio.get_running_loop()
return self

async def __aexit__(self, *_dont_care: object) -> None:
Expand Down

0 comments on commit 2c0b8ad

Please sign in to comment.