NEW: generator hooks ✨

### Motivation

I need to add some custom metric based on when exactly the decorated function got called: a background job start delay after its scheduled execution time.

### Implementation

In addition to simple functions, the decorator now accepts generator functions. That would allow them to run some code just before the decorated function gets called and remember the state.

Plain old hooks are not affected. The change is designed to be backwards-compatible.

pyproject.toml

@@ -28,6 +28,9 @@ omit = ["tests/*"]
 [build-system]
 requires = ["setuptools>=42", "wheel", "setuptools_scm[toml]>=3.4"]
 
+[tool.mypy]
+warn_unused_configs = true
+
 [[tool.mypy.overrides]]
 module = ["mock", "freezegun", "elasticsearch.*", "fqn_decorators.*", "pkgsettings", "setuptools", "Queue"]
 ignore_missing_imports = true

tests/test_hooks.py

@@ -2,7 +2,7 @@
 from fqn_decorators import get_fqn
 
 from tests.conftest import go
-from time_execution import settings, time_execution
+from time_execution import GeneratorHookReturnType, settings, time_execution
 from time_execution.backends.base import BaseMetricsBackend
 
 
@@ -231,3 +231,34 @@ def hook(response, exception, metric, func, func_args, func_kwargs):
 
         with settings(hooks=[hook]):
             go(param1=param)
+
+    def test_generator_hook(self) -> None:
+        def generator_hook(func, func_args, func_kwargs) -> GeneratorHookReturnType:
+            assert func_args == (42,)
+            assert func_kwargs == {"bar": 100500}
+            (response, _exception, _metrics) = yield
+            assert response == "response"
+            return {"key": "value"}
+
+        @time_execution(disable_default_hooks=True, extra_hooks=(generator_hook,))
+        def go(foo: int, *, bar: int) -> str:
+            return "response"
+
+        def asserts(_name, **data):
+            assert data["key"] == "value"
+
+        with settings(backends=[AssertBackend(asserts)]):
+            assert go(42, bar=100500) == "response"
+
+    def test_generator_hook_did_not_stop(self) -> None:
+        def generator_hook(func, func_args, func_kwargs) -> GeneratorHookReturnType:
+            yield
+            yield  # this extra `yield` is incorrect
+            return {}
+
+        @time_execution(disable_default_hooks=True, extra_hooks=(generator_hook,))
+        def go() -> None:
+            return None
+
+        with pytest.raises(RuntimeError, match="generator hook did not stop"):
+            go()

time_execution/decorator.py

@@ -5,11 +5,11 @@
 from asyncio import iscoroutinefunction
 from collections.abc import Iterable
 from functools import wraps
-from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, cast
+from typing import Any, Callable, Dict, Generator, Optional, Tuple, TypeVar, cast
 
 import fqn_decorators
 from pkgsettings import Settings
-from typing_extensions import Protocol, overload
+from typing_extensions import Protocol, TypeAlias, overload
 
 _F = TypeVar("_F", bound=Callable[..., Any])
 
@@ -31,7 +31,7 @@ def time_execution(__wrapped: _F) -> _F:
 def time_execution(
     *,
     get_fqn: Callable[[Any], str] = fqn_decorators.get_fqn,
-    extra_hooks: Optional[Iterable[Hook]] = None,
+    extra_hooks: Optional[Iterable[Hook | GeneratorHook]] = None,
     disable_default_hooks: bool = False,
 ) -> Callable[[_F], _F]:
     """
@@ -91,3 +91,25 @@ def __call__(
         func_kwargs: Dict[str, Any],
     ) -> Optional[Dict[str, Any]]:
         ...
+
+
+GeneratorHookReturnType: TypeAlias = Generator[
+    None, Tuple[Any, Optional[BaseException], Dict[str, Any]], Optional[Dict[str, Any]]
+]
+
+
+class GeneratorHook(Protocol):
+    """
+    Generator-type hook.
+
+    This kind of hook gets called before the target function.
+    Response, exception, and metrics are sent into the generator after the target function finishes.
+    """
+
+    def __call__(
+        self,
+        func: Callable[..., Any],
+        func_args: Tuple[Any, ...],
+        func_kwargs: Dict[str, Any],
+    ) -> GeneratorHookReturnType:
+        ...

time_execution/timed.py

@@ -2,12 +2,13 @@
 
 from collections.abc import Iterable
 from contextlib import AbstractContextManager
+from inspect import isgenerator, isgeneratorfunction
 from socket import gethostname
 from timeit import default_timer
 from types import TracebackType
-from typing import Any, Callable, Dict, Optional, Tuple, Type
+from typing import Any, Callable, Dict, Optional, Tuple, Type, cast
 
-from time_execution import Hook, settings, write_metric
+from time_execution import GeneratorHook, GeneratorHookReturnType, Hook, settings, write_metric
 
 SHORT_HOSTNAME = gethostname()
 
@@ -22,8 +23,7 @@ class Timed(AbstractContextManager):
         "result",
         "_wrapped",
         "_fqn",
-        "_extra_hooks",
-        "_disable_default_hooks",
+        "_hooks",
         "_call_args",
         "_call_kwargs",
         "_start_time",
@@ -36,19 +36,31 @@ def __init__(
         fqn: str,
         call_args: Tuple[Any, ...],
         call_kwargs: Dict[str, Any],
-        extra_hooks: Optional[Iterable[Hook]] = None,
+        extra_hooks: Optional[Iterable[Hook | GeneratorHook]] = None,
         disable_default_hooks: bool = False,
     ) -> None:
         self.result: Optional[Any] = None
         self._wrapped = wrapped
         self._fqn = fqn
-        self._extra_hooks = extra_hooks
-        self._disable_default_hooks = disable_default_hooks
         self._call_args = call_args
         self._call_kwargs = call_kwargs
 
+        hooks = extra_hooks or ()
+        if not disable_default_hooks:
+            hooks = (*settings.hooks, *hooks)
+
+        self._hooks = tuple(
+            cast(Hook, hook) if not isgeneratorfunction(hook)  # simple hook, we'll call it in the exit
+            # For a generator hook, call it. We'll start in the entrance.
+            else cast(GeneratorHookReturnType, hook(func=wrapped, func_args=call_args, func_kwargs=call_kwargs))
+            for hook in hooks
+        )
+
     def __enter__(self) -> Timed:
         self._start_time = default_timer()
+        for hook in self._hooks:
+            if isgenerator(hook):
+                hook.send(None)  # start a generator hook
         return self
 
     def __exit__(
@@ -65,14 +77,9 @@ def __exit__(
         if origin:
             metric["origin"] = origin
 
-        hooks = self._extra_hooks or ()
-        if not self._disable_default_hooks:
-            hooks = (*settings.hooks, *hooks)
-
         # Apply the registered hooks, and collect the metadata they might
         # return to be stored with the metrics.
         metadata = self._apply_hooks(
-            hooks=hooks,
             response=self.result,
             exception=__exc_val,
             metric=metric,
@@ -81,17 +88,27 @@ def __exit__(
         metric.update(metadata)
         write_metric(**metric)  # type: ignore[arg-type]
 
-    def _apply_hooks(self, hooks, response, exception, metric) -> Dict:
-        metadata = dict()
-        for hook in hooks:
-            hook_result = hook(
-                response=response,
-                exception=exception,
-                metric=metric,
-                func=self._wrapped,
-                func_args=self._call_args,
-                func_kwargs=self._call_kwargs,
-            )
+    def _apply_hooks(self, response, exception, metric) -> Dict:
+        metadata: Dict[str, Any] = dict()
+        for hook in self._hooks:
+            if not isgenerator(hook):
+                # Simple exit hook, call it directly.
+                hook_result = cast(Hook, hook)(
+                    response=response,
+                    exception=exception,
+                    metric=metric,
+                    func=self._wrapped,
+                    func_args=self._call_args,
+                    func_kwargs=self._call_kwargs,
+                )
+            else:
+                # Generator hook: send the results and obtain custom metadata.
+                try:
+                    hook.send((response, exception, metric))
+                except StopIteration as e:
+                    hook_result = e.value
+                else:
+                    raise RuntimeError("generator hook did not stop")
             if hook_result:
                 metadata.update(hook_result)
         return metadata