From 29f42e38fc0aef76af67f7d4725d90a7903f9d2c Mon Sep 17 00:00:00 2001 From: Sergey Vasilyev Date: Wed, 24 Dec 2025 16:06:00 +0100 Subject: [PATCH] Add the loop time enabler for explicit usage (both as a context manager and decorator) Signed-off-by: Sergey Vasilyev --- looptime/__init__.py | 2 + looptime/enabler.py | 105 ++++++++++++++++++++++++++++++++++++++++++ tests/test_enabler.py | 79 +++++++++++++++++++++++++++++++ 3 files changed, 186 insertions(+) create mode 100644 looptime/enabler.py create mode 100644 tests/test_enabler.py diff --git a/looptime/__init__.py b/looptime/__init__.py index 3882114..d25a75c 100644 --- a/looptime/__init__.py +++ b/looptime/__init__.py @@ -1,4 +1,5 @@ from .chronometers import Chronometer +from .enabler import enabled from .loops import IdleTimeoutError, LoopTimeEventLoop, LoopTimeoutError, TimeWarning from .patchers import make_event_loop_class, new_event_loop, patch_event_loop, reset_caches from .timeproxies import LoopTimeProxy @@ -14,4 +15,5 @@ 'new_event_loop', 'patch_event_loop', 'make_event_loop_class', + 'enabled', ] diff --git a/looptime/enabler.py b/looptime/enabler.py new file mode 100644 index 0000000..cf84d07 --- /dev/null +++ b/looptime/enabler.py @@ -0,0 +1,105 @@ +import asyncio +import functools +import inspect +import warnings +from typing import Any, Callable, ContextManager, ParamSpec, TypeVar + +from looptime import loops + +P = ParamSpec('P') +R = TypeVar('R') + + +class enabled(ContextManager[None]): + """ + Enable the looptime time compaction temporarily. + + If used as a context manager, enables the time compaction for the wrapped + code block only:: + + import asyncio + import looptime + + async def main() -> None: + with looptime.enabled(strict=True): + await asyncio.sleep(10) + + if __name__ == '__main__': + asuncio.run(main()) + + If used as a function/fixture decorator, enables the time compaction + for the duration of the function/fixture:: + + import asyncio + import looptime + + @looptime.enabled(strict=True) + async def main() -> None: + await asyncio.sleep(10) + + if __name__ == '__main__': + asuncio.run(main()) + + In both cases, the event loop must be pre-patched (usually at creation). + In strict mode, if the event loop is not patched, the call will fail. + In non-strict mode (the default), it will issue a warning and continue + with the real time flow (i.e. with no time compaction). + + Use it, for example, for fixtures or finalizers of fixtures where the fast + time flow is required despite fixtures are normally excluded from the time + compaction magic (because it is impossible or difficult to infer which + event loop is being used in the multi-scoped setup of pytest-asyncio), + and because of the structure of pytest hooks for fixture finalizing + (no finalizer hook, only the post-finalizer hook, when it is too late). + + Beware of a caveat: if used as a decorator on a yield-based fixture, + it will enable the looptime magic for the whole duration of the test, + including all its fixtures (even undecorated ones), until the decorated + fixture reaches its finalizer. This might have unexpected side effects. + """ + strict: bool + _loop: asyncio.AbstractEventLoop | None + _mgr: ContextManager[None] | None + + def __init__(self, *, strict: bool = False, loop: asyncio.AbstractEventLoop | None = None) -> None: + super().__init__() + self.strict = strict + self._loop = loop + self._mgr = None + + def __enter__(self) -> None: + msg = "The running loop is not a looptime-patched loop, cannot enable it." + loop = self._loop if self._loop is not None else asyncio.get_running_loop() + if isinstance(loop, loops.LoopTimeEventLoop): + self._mgr = loop.looptime_enabled() + self._mgr.__enter__() + elif self.strict: + raise RuntimeError(msg) + else: + warnings.warn(msg, UserWarning) + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + if self._mgr is not None: + self._mgr.__exit__(exc_type, exc_val, exc_tb) + self._mgr = None + + # Type checkers: too complicated. We get R=Coroutine[Y,S,RR] for async functions, + # but return that last RR part, which turns to be Any. The runtime is unaffected. + # I don't know how to properly annotate such a mixed sync-async decorator internally. + # The external declaration of __call__() is sufficient and correct. + # TODO: LATER: try annotating it properly. + def __call__(self, fn: Callable[P, R]) -> Callable[P, R]: + if inspect.iscoroutinefunction(fn): + @functools.wraps(fn) + async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + nonlocal self + with self: + return await fn(*args, **kwargs) # type: ignore + else: + @functools.wraps(fn) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + nonlocal self + with self: + return fn(*args, **kwargs) + + return wrapper # type: ignore diff --git a/tests/test_enabler.py b/tests/test_enabler.py new file mode 100644 index 0000000..95881bc --- /dev/null +++ b/tests/test_enabler.py @@ -0,0 +1,79 @@ +import asyncio +import sys + +import pytest + +import looptime + + +@pytest.mark.asyncio +async def test_enabler_as_context_manager(): + loop = asyncio.get_running_loop() + enabled = isinstance(loop, looptime.LoopTimeEventLoop) and loop.looptime_on + assert not enabled + + with looptime.enabled(): + enabled = isinstance(loop, looptime.LoopTimeEventLoop) and loop.looptime_on + assert enabled + + +@pytest.mark.asyncio +async def test_enabler_as_decorator_for_sync_functions(): + @looptime.enabled() + def fn(a: int) -> tuple[int, bool]: + loop = asyncio.get_running_loop() + enabled = isinstance(loop, looptime.LoopTimeEventLoop) and loop.looptime_on + return a + 10, enabled + + loop = asyncio.get_running_loop() + enabled = isinstance(loop, looptime.LoopTimeEventLoop) and loop.looptime_on + assert not enabled + + result, enabled = fn(123) + assert result == 133 + assert enabled + + +@pytest.mark.asyncio +async def test_enabler_as_decorator_for_async_functions(): + @looptime.enabled() + async def fn(a: int) -> tuple[int, bool]: + loop = asyncio.get_running_loop() + enabled = isinstance(loop, looptime.LoopTimeEventLoop) and loop.looptime_on + return a + 10, enabled + + loop = asyncio.get_running_loop() + enabled = isinstance(loop, looptime.LoopTimeEventLoop) and loop.looptime_on + assert not enabled + + result, enabled = await fn(123) + assert result == 133 + assert enabled + + +@pytest.mark.skipif(sys.version_info < (3, 11), reason="Runners require Python>=3.11") +def test_enabler_with_explicit_loop(): + with asyncio.Runner() as runner: + runner_loop = runner.get_loop() + looptime.patch_event_loop(runner_loop, _enabled=False) + with looptime.enabled(loop=runner_loop): + enabled = isinstance(runner_loop, looptime.LoopTimeEventLoop) and runner_loop.looptime_on + assert enabled + + +@pytest.mark.skipif(sys.version_info < (3, 11), reason="Runners require Python>=3.11") +def test_strict_mode_error(): + with asyncio.Runner() as runner: + runner_loop = runner.get_loop() # unpatched! + with pytest.raises(RuntimeError, match="loop is not a looptime-patched loop"): + with looptime.enabled(loop=runner_loop, strict=True): + pass + + +@pytest.mark.skipif(sys.version_info < (3, 11), reason="Runners require Python>=3.11") +def test_nonstrict_mode_warning(): + with asyncio.Runner() as runner: + runner_loop = runner.get_loop() # unpatched! + with pytest.warns(UserWarning, match="loop is not a looptime-patched loop"): + with looptime.enabled(loop=runner_loop, strict=False): + pass