dycw-utilities 0.121.1__py3-none-any.whl → 0.122.1__py3-none-any.whl

{dycw_utilities-0.121.1.dist-info → dycw_utilities-0.122.1.dist-info}/METADATA

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

{dycw_utilities-0.121.1.dist-info → dycw_utilities-0.122.1.dist-info}/RECORD

@@ -1,7 +1,7 @@
-utilities/__init__.py,sha256=Ses6BHZNi0jOdocz8doinXp-udEMKae9OJkWZ9Ptqww,60
+utilities/__init__.py,sha256=HQ1CqsUZ8X8-X0sFT3wvjjj0NMYqAnzvjRBDmvWB2xU,60
 utilities/altair.py,sha256=Gpja-flOo-Db0PIPJLJsgzAlXWoKUjPU1qY-DQ829ek,9156
 utilities/astor.py,sha256=xuDUkjq0-b6fhtwjhbnebzbqQZAjMSHR1IIS5uOodVg,777
-utilities/asyncio.py,sha256=gn8dxBsNzFfFzrNpImF731LAmgLwEQM3uiwKuUPlg08,18822
+utilities/asyncio.py,sha256=cjXpSezv95ZG3t-D9sXLOn6BGBj7w-Wf4udvE5gU4vU,23508
 utilities/atomicwrites.py,sha256=geFjn9Pwn-tTrtoGjDDxWli9NqbYfy3gGL6ZBctiqSo,5393
 utilities/atools.py,sha256=IYMuFSFGSKyuQmqD6v5IUtDlz8PPw0Sr87Cub_gRU3M,1168
 utilities/cachetools.py,sha256=C1zqOg7BYz0IfQFK8e3qaDDgEZxDpo47F15RTfJM37Q,2910
@@ -16,7 +16,7 @@ utilities/datetime.py,sha256=uYoaOi_C1YtNXGfTN9xlTrW62Re2b1_4Skuv14_MeYQ,38985
 utilities/enum.py,sha256=HoRwVCWzsnH0vpO9ZEcAAIZLMv0Sn2vJxxA4sYMQgDs,5793
 utilities/errors.py,sha256=gxsaa7eq7jbYl41Of40-ivjXqJB5gt4QAcJ0smZZMJE,829
 utilities/eventkit.py,sha256=6M5Xu1SzN-juk9PqBHwy5dS-ta7T0qA6SMpDsakOJ0E,13039
-utilities/fastapi.py,sha256=eiisloI6kQVCkPfDpBzlLrDZDi8yJ0VmrSPlJ2k84Mo,2334
+utilities/fastapi.py,sha256=LG1-Q8RDi7wsyVN6v74qptPYX8WGXPkFOQFniMvtzjc,2439
 utilities/fpdf2.py,sha256=y1NGXR5chWqLXWpewGV3hlRGMr_5yV1lVRkPBhPEgJI,1843
 utilities/functions.py,sha256=jgt592voaHNtX56qX0SRvFveVCRmSIxCZmqvpLZCnY8,27305
 utilities/functools.py,sha256=WrpHt7NLNWSUn9A1Q_ZIWlNaYZOEI4IFKyBG9HO3BC4,1643
@@ -59,7 +59,7 @@ utilities/pytest_regressions.py,sha256=-SVT9647Dg6-JcdsiaDKXe3NdOmmrvGevLKWwGjxq
 utilities/python_dotenv.py,sha256=iWcnpXbH7S6RoXHiLlGgyuH6udCupAcPd_gQ0eAenQ0,3190
 utilities/random.py,sha256=lYdjgxB7GCfU_fwFVl5U-BIM_HV3q6_urL9byjrwDM8,4157
 utilities/re.py,sha256=5J4d8VwIPFVrX2Eb8zfoxImDv7IwiN_U7mJ07wR2Wvs,3958
-utilities/redis.py,sha256=i5G3k-EIi2JSf9arLIu0YW6viGsiYtxAoynbb4KvJek,26628
+utilities/redis.py,sha256=8ELnXiISVHY1m9elJWhekhLXg6NQSaNIIPvZ_EaGw3s,26624
 utilities/reprlib.py,sha256=Re9bk3n-kC__9DxQmRlevqFA86pE6TtVfWjUgpbVOv0,1849
 utilities/rich.py,sha256=t50MwwVBsoOLxzmeVFSVpjno4OW6Ufum32skXbV8-Bs,1911
 utilities/scipy.py,sha256=X6ROnHwiUhAmPhM0jkfEh0-Fd9iRvwiqtCQMOLmOQF8,945
@@ -67,7 +67,7 @@ utilities/sentinel.py,sha256=3jIwgpMekWgDAxPDA_hXMP2St43cPhciKN3LWiZ7kv0,1248
 utilities/shelve.py,sha256=HZsMwK4tcIfg3sh0gApx4-yjQnrY4o3V3ZRimvRhoW0,738
 utilities/slack_sdk.py,sha256=A2f7-DYOngRoUP6ZdLIaUQ6Lfzgru5Xp3U3k5JfEkQE,3301
 utilities/socket.py,sha256=K77vfREvzoVTrpYKo6MZakol0EYu2q1sWJnnZqL0So0,118
-utilities/sqlalchemy.py,sha256=iwNJyR0Dx47l81aHEq5P6eMzgduVxxGH531OVUp2k_A,35451
+utilities/sqlalchemy.py,sha256=yCUCDhg0wFOCdEh6wwBD7Ma979OksLpz4arck1s653s,35447
 utilities/sqlalchemy_polars.py,sha256=wjJpoUo-yO9E2ujpG_06vV5r2OdvBiQ4yvV6wKCa2Tk,15605
 utilities/statsmodels.py,sha256=koyiBHvpMcSiBfh99wFUfSggLNx7cuAw3rwyfAhoKpQ,3410
 utilities/streamlit.py,sha256=U9PJBaKP1IdSykKhPZhIzSPTZsmLsnwbEPZWzNhJPKk,2955
@@ -88,7 +88,7 @@ utilities/warnings.py,sha256=un1LvHv70PU-LLv8RxPVmugTzDJkkGXRMZTE2-fTQHw,1771
 utilities/whenever.py,sha256=jS31ZAY5OMxFxLja_Yo5Fidi87Pd-GoVZ7Vi_teqVDA,16743
 utilities/zipfile.py,sha256=24lQc9ATcJxHXBPc_tBDiJk48pWyRrlxO2fIsFxU0A8,699
 utilities/zoneinfo.py,sha256=-5j7IQ9nb7gR43rdgA7ms05im-XuqhAk9EJnQBXxCoQ,1874
-dycw_utilities-0.121.1.dist-info/METADATA,sha256=54NwpwmQD2-MOcByhCWt9vZDiig3pdJSvdA5VmjhNbY,12943
-dycw_utilities-0.121.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.121.1.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.121.1.dist-info/RECORD,,
+dycw_utilities-0.122.1.dist-info/METADATA,sha256=9QN9Ulm7O-TsU1Zg2xrl5ylNadZMkHMp9ktRVFxpFwQ,12943
+dycw_utilities-0.122.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.122.1.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.122.1.dist-info/RECORD,,

utilities/__init__.py

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

utilities/asyncio.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import datetime as dt
 from abc import ABC, abstractmethod
 from asyncio import (
+    CancelledError,
     Event,
     PriorityQueue,
     Queue,
@@ -12,11 +13,18 @@ from asyncio import (
     Task,
     TaskGroup,
     create_subprocess_shell,
+    create_task,
     sleep,
     timeout,
 )
 from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping
-from contextlib import _AsyncGeneratorContextManager, asynccontextmanager
+from contextlib import (
+    AbstractAsyncContextManager,
+    AsyncExitStack,
+    _AsyncGeneratorContextManager,
+    asynccontextmanager,
+    suppress,
+)
 from dataclasses import dataclass, field
 from io import StringIO
 from logging import getLogger
@@ -27,6 +35,7 @@ from typing import (
     Any,
     Generic,
     NoReturn,
+    Self,
     TextIO,
     TypeVar,
     assert_never,
@@ -42,7 +51,7 @@ from utilities.datetime import (
     get_now,
     round_datetime,
 )
-from utilities.errors import repr_error
+from utilities.errors import ImpossibleCaseError, repr_error
 from utilities.functions import ensure_int, ensure_not_none, get_class_name
 from utilities.reprlib import get_repr
 from utilities.sentinel import Sentinel, sentinel
@@ -60,6 +69,7 @@ if TYPE_CHECKING:
     from asyncio.subprocess import Process
     from collections.abc import AsyncIterator, Sequence
     from contextvars import Context
+    from types import TracebackType
 
     from utilities.types import Duration
 
@@ -76,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
@@ -90,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,
@@ -107,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:
@@ -125,26 +157,94 @@ class EnhancedTaskGroup(TaskGroup):
 class InfiniteLooper(ABC, Generic[THashable]):
     """An infinite loop which can throw exceptions by setting events."""
 
-    sleep_core: DurationOrEveryDuration = SECOND
-    sleep_restart: DurationOrEveryDuration = MINUTE
-    logger: str | None = None
+    sleep_core: DurationOrEveryDuration = field(default=SECOND, repr=False)
+    sleep_restart: DurationOrEveryDuration = field(default=MINUTE, repr=False)
+    duration: Duration | None = field(default=None, repr=False)
+    logger: str | None = field(default=None, repr=False)
+    _await_upon_aenter: bool = field(default=True, init=False, repr=False)
+    _depth: int = field(default=0, init=False, repr=False)
     _events: Mapping[THashable | None, Event] = field(
         default_factory=dict, init=False, repr=False, hash=False
     )
+    _stack: AsyncExitStack = field(
+        default_factory=AsyncExitStack, init=False, repr=False
+    )
+    _task: Task[None] | None = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
         self._events = {
             event: Event() for event, _ in self._yield_events_and_exceptions()
         }
 
-    async def __call__(self) -> None:
-        """Create a coroutine to run the looper."""
-        coroutines = list(self._yield_coroutines())
-        if len(coroutines) == 0:
-            return await self._run_looper()
-        return await self._run_looper_with_coroutines(*coroutines)
+    async def __aenter__(self) -> Self:
+        """Context manager entry."""
+        if (self._task is None) and (self._depth == 0):
+            _ = await self._stack.__aenter__()
+            self._task = create_task(self._run_looper())
+            if self._await_upon_aenter:
+                with suppress(CancelledError):
+                    await self._task
+        elif (self._task is not None) and (self._depth >= 1):
+            ...
+        else:
+            raise ImpossibleCaseError(  # pragma: no cover
+                case=[f"{self._task=}", f"{self._depth=}"]
+            )
+        self._depth += 1
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        """Context manager exit."""
+        _ = (exc_type, exc_value, traceback)
+        if (self._task is None) or (self._depth == 0):
+            raise ImpossibleCaseError(  # pragma: no cover
+                case=[f"{self._task=}", f"{self._depth=}"]
+            )
+        self._depth -= 1
+        if self._depth == 0:
+            _ = await self._stack.__aexit__(exc_type, exc_value, traceback)
+            with suppress(CancelledError):
+                await self._task
+            self._task = None
+            with suppress(Exception):
+                await self._teardown()
+
+    async def stop(self) -> None:
+        """Stop the service."""
+        if self._task is None:
+            raise ImpossibleCaseError(case=[f"{self._task=}"])  # pragma: no cover
+        with suppress(CancelledError):
+            _ = self._task.cancel()
 
     async def _run_looper(self) -> None:
+        """Run the looper."""
+        match self.duration:
+            case None:
+                await self._run_looper_without_timeout()
+            case int() | float() | dt.timedelta() as duration:
+                try:
+                    async with timeout_dur(duration=duration):
+                        return await self._run_looper_without_timeout()
+                except TimeoutError:
+                    await self.stop()
+            case _ as never:
+                assert_never(never)
+        return None
+
+    async def _run_looper_without_timeout(self) -> None:
+        """Run the looper without a timeout."""
+        coroutines = list(self._yield_coroutines())
+        loopers = list(self._yield_loopers())
+        if (len(coroutines) == 0) and (len(loopers) == 0):
+            return await self._run_looper_by_itself()
+        return await self._run_looper_with_others(coroutines, loopers)
+
+    async def _run_looper_by_itself(self) -> None:
         """Run the looper by itself."""
         while True:
             try:
@@ -171,20 +271,31 @@ class InfiniteLooper(ABC, Generic[THashable]):
                 raise
             except Exception as error:  # noqa: BLE001
                 self._error_upon_core(error)
-                await self._run_sleep(self.sleep_restart)
+                try:
+                    await self._teardown()
+                except Exception as error:  # noqa: BLE001
+                    self._error_upon_teardown(error)
+                finally:
+                    await self._run_sleep(self.sleep_restart)
 
-    async def _run_looper_with_coroutines(
-        self, *coroutines: Callable[[], Coroutine1[None]]
+    async def _run_looper_with_others(
+        self,
+        coroutines: Iterable[Callable[[], Coroutine1[None]]],
+        loopers: Iterable[InfiniteLooper[Any]],
+        /,
     ) -> None:
         """Run multiple loopers."""
         while True:
             self._reset_events()
             try:
-                async with TaskGroup() as tg:
-                    _ = tg.create_task(self._run_looper())
+                async with TaskGroup() as tg, AsyncExitStack() as stack:
+                    _ = tg.create_task(self._run_looper_by_itself())
                     _ = [tg.create_task(c()) for c in coroutines]
+                    _ = [
+                        tg.create_task(stack.enter_async_context(lo)) for lo in loopers
+                    ]
             except ExceptionGroup as error:
-                self._error_group_upon_coroutines(error)
+                self._error_group_upon_others(error)
                 await self._run_sleep(self.sleep_restart)
 
     async def _initialize(self) -> None:
@@ -193,6 +304,9 @@ class InfiniteLooper(ABC, Generic[THashable]):
     async def _core(self) -> None:
         """Run the core part of the loop."""
 
+    async def _teardown(self) -> None:
+        """Tear down the loop."""
+
     def _error_upon_initialize(self, error: Exception, /) -> None:
         """Handle any errors upon initializing the looper."""
         if self.logger is not None:
@@ -213,7 +327,17 @@ class InfiniteLooper(ABC, Generic[THashable]):
                 self._sleep_restart_desc,
             )
 
-    def _error_group_upon_coroutines(self, group: ExceptionGroup, /) -> None:
+    def _error_upon_teardown(self, error: Exception, /) -> None:
+        """Handle any errors upon tearing down the looper."""
+        if self.logger is not None:
+            getLogger(name=self.logger).error(
+                "%r encountered %r whilst tearing down; sleeping %s...",
+                get_class_name(self),
+                repr_error(error),
+                self._sleep_restart_desc,
+            )
+
+    def _error_group_upon_others(self, group: ExceptionGroup, /) -> None:
         """Handle any errors upon running the core function."""
         if self.logger is not None:
             errors = group.exceptions
@@ -269,15 +393,19 @@ class InfiniteLooper(ABC, Generic[THashable]):
             raise _InfiniteLooperNoSuchEventError(looper=self, event=event) from None
         event_obj.set()
 
+    def _yield_events_and_exceptions(
+        self,
+    ) -> Iterator[tuple[THashable | None, MaybeType[Exception]]]:
+        """Yield the events & exceptions."""
+        yield (None, _InfiniteLooperDefaultEventError(looper=self))
+
     def _yield_coroutines(self) -> Iterator[Callable[[], Coroutine1[None]]]:
         """Yield any other coroutines which must also be run."""
         yield from []
 
-    def _yield_events_and_exceptions(
-        self,
-    ) -> Iterator[tuple[THashable | None, MaybeType[BaseException]]]:
-        """Yield the events & exceptions."""
-        yield (None, _InfiniteLooperDefaultEventError)
+    def _yield_loopers(self) -> Iterator[InfiniteLooper[Any]]:
+        """Yield any other loopers which must also be run."""
+        yield from []
 
 
 @dataclass(kw_only=True, slots=True)
@@ -309,6 +437,7 @@ class InfiniteQueueLooper(InfiniteLooper[THashable], Generic[THashable, _T]):
     """An infinite loop which processes a queue."""
 
     queue_type: type[Queue[_T]] = field(default=Queue, repr=False)
+    _await_upon_aenter: bool = field(default=False, init=False, repr=False)
     _queue: Queue[_T] = field(init=False)
 
     @override
@@ -346,6 +475,7 @@ class InfiniteQueueLooper(InfiniteLooper[THashable], Generic[THashable, _T]):
         """Run until the queue is empty."""
         while not self.empty():
             await self._process_items(*get_items_nowait(self._queue))
+        await self.stop()
 
     @override
     def _error_upon_core(self, error: Exception, /) -> None:
@@ -353,12 +483,12 @@ class InfiniteQueueLooper(InfiniteLooper[THashable], Generic[THashable, _T]):
         if self.logger is not None:
             if isinstance(error, InfiniteQueueLooperError):
                 getLogger(name=self.logger).error(
-                    "%r encountered %s whilst processing %d item(s) %s; sleeping for %s...",
+                    "%r encountered %s whilst processing %d item(s) %s; sleeping %s...",
                     get_class_name(self),
                     repr_error(error.error),
                     len(error.items),
                     get_repr(error.items),
-                    self.sleep_restart,
+                    self._sleep_restart_desc,
                 )
             else:
                 super()._error_upon_core(error)  # pragma: no cover
@@ -370,10 +500,6 @@ class InfiniteQueueLooperError(Exception, Generic[_T]):
     items: Sequence[_T]
     error: Exception
 
-    @override
-    def __str__(self) -> str:
-        return f"{get_class_name(self.looper)!r} encountered {repr_error(self.error)} whilst processing {len(self.items)} item(s): {get_repr(self.items)}"
-
 
 ##
 

utilities/fastapi.py

@@ -70,5 +70,9 @@ class PingReceiver(InfiniteLooper):
     async def _initialize(self) -> None:
         await self._server.serve()  # skipif-ci
 
+    @override
+    async def _teardown(self) -> None:
+        await self._server.shutdown()  # skipif-ci
+
 
 __all__ = ["PingReceiver"]

utilities/redis.py

@@ -611,7 +611,7 @@ class Publisher(InfiniteQueueLooper[None, tuple[str, EncodableT]]):
     @override
     def _yield_events_and_exceptions(
         self,
-    ) -> Iterator[tuple[None, MaybeType[BaseException]]]:
+    ) -> Iterator[tuple[None, MaybeType[Exception]]]:
         yield (None, PublisherError)  # skipif-ci-and-not-linux
 
 

utilities/sqlalchemy.py

@@ -640,7 +640,7 @@ class Upserter(InfiniteQueueLooper[None, _InsertItem]):
     @override
     def _yield_events_and_exceptions(
         self,
-    ) -> Iterator[tuple[None, MaybeType[BaseException]]]:
+    ) -> Iterator[tuple[None, MaybeType[Exception]]]:
         yield (None, UpserterError)