dycw-utilities 0.122.0__tar.gz → 0.122.1__tar.gz

{dycw_utilities-0.122.0 → dycw_utilities-0.122.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.122.0
+Version: 0.122.1
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.122.0 → dycw_utilities-0.122.1}/pyproject.toml

@@ -92,7 +92,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.122.0"
+version = "0.122.1"
 
 [project.optional-dependencies]
 test = [
@@ -334,7 +334,7 @@ zzz-test-zoneinfo = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.122.0"
+current_version = "0.122.1"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.122.0 → dycw_utilities-0.122.1}/src/tests/test_asyncio.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from asyncio import CancelledError, Event, Queue, TaskGroup, run, sleep, timeout
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from functools import partial
 from itertools import chain, count
@@ -53,7 +54,7 @@ from utilities.sentinel import Sentinel, sentinel
 from utilities.timer import Timer
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterator
+    from collections.abc import AsyncIterator, Callable, Iterator
 
     from utilities.types import (
         Coroutine1,
@@ -65,6 +66,41 @@ if TYPE_CHECKING:
 
 
 class TestEnhancedTaskGroup:
+    async def test_context_managers(self) -> None:
+        first: bool = False
+
+        @asynccontextmanager
+        async def yield_first() -> AsyncIterator[None]:
+            nonlocal first
+            try:
+                first = True
+                yield
+            finally:
+                first = False
+
+        second: bool = False
+
+        @asynccontextmanager
+        async def yield_second() -> AsyncIterator[None]:
+            nonlocal second
+            try:
+                second = True
+                yield
+            finally:
+                second = False
+
+        async with EnhancedTaskGroup() as tg:
+            assert not first
+            assert not second
+            _ = await tg.enter_async_context(yield_first())
+            assert first
+            assert not second
+            _ = await tg.enter_async_context(yield_second())
+            assert first
+            assert second
+        assert not first
+        assert not second
+
     async def test_max_tasks_disabled(self) -> None:
         with Timer() as timer:
             async with EnhancedTaskGroup() as tg:
@@ -346,7 +382,7 @@ class TestInfiniteLooper:
             ...
         assert 4 <= looper.initializations <= 6
         assert 0 <= looper.counter <= 7
-        assert 16 <= external <= 21
+        assert 14 <= external <= 21
 
     async def test_with_coroutine_self_error(self) -> None:
         class CustomError(Exception): ...

{dycw_utilities-0.122.0 → dycw_utilities-0.122.1}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.122.0"
+__version__ = "0.122.1"

{dycw_utilities-0.122.0 → dycw_utilities-0.122.1}/src/utilities/asyncio.py

@@ -19,6 +19,7 @@ from asyncio import (
 )
 from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping
 from contextlib import (
+    AbstractAsyncContextManager,
     AsyncExitStack,
     _AsyncGeneratorContextManager,
     asynccontextmanager,
@@ -85,6 +86,8 @@ class EnhancedTaskGroup(TaskGroup):
     _semaphore: Semaphore | None
     _timeout: Duration | None
     _error: type[Exception]
+    _stack: AsyncExitStack
+    _stack_entered: bool
     _timeout_cm: _AsyncGeneratorContextManager[None] | None
 
     @override
@@ -99,8 +102,25 @@ class EnhancedTaskGroup(TaskGroup):
         self._semaphore = None if max_tasks is None else Semaphore(max_tasks)
         self._timeout = timeout
         self._error = error
+        self._stack = AsyncExitStack()
+        self._stack_entered = False  # TOOD: no need
         self._timeout_cm = None
 
+    @override
+    async def __aenter__(self) -> Self:
+        _ = await self._stack.__aenter__()
+        return await super().__aenter__()
+
+    @override
+    async def __aexit__(
+        self,
+        et: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        _ = await self._stack.__aexit__(et, exc, tb)
+        _ = await super().__aexit__(et, exc, tb)
+
     @override
     def create_task(
         self,
@@ -116,6 +136,9 @@ class EnhancedTaskGroup(TaskGroup):
         coroutine = self._wrap_with_timeout(coroutine)
         return super().create_task(coroutine, name=name, context=context)
 
+    async def enter_async_context(self, cm: AbstractAsyncContextManager[_T], /) -> _T:
+        return await self._stack.enter_async_context(cm)
+
     async def _wrap_with_semaphore(
         self, semaphore: Semaphore, coroutine: _CoroutineLike[_T], /
     ) -> _T: