Add the loop time enabler for explicit usage (both as a context manager and decorator)

Signed-off-by: Sergey Vasilyev <nolar@nolar.info>

Author: nolar

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',
 ]

looptime/enabler.py

@@ -0,0 +1,100 @@
+import asyncio
+import functools
+import inspect
+import warnings
+from typing import 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, exc_val, exc_tb) -> None:
+        if self._mgr is not None:
+            self._mgr.__exit__(exc_type, exc_val, exc_tb)
+            self._mgr = None
+
+    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)
+        else:
+            @functools.wraps(fn)
+            def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                nonlocal self
+                with self:
+                    return fn(*args, **kwargs)
+
+        return wrapper

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))  # for Runners
+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))  # for Runners
+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))  # for Runners
+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