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.

README.md

@@ -176,13 +176,16 @@ With a hook you can add additional and change existing fields. This can
 be useful for cases where you would like to add a column to the metric
 based on the response of the wrapped function.
 
-A hook will always get 3 arguments:
+### Simple callable hook
 
-* `response` - The returned value of the wrapped function
-* `exception` - The raised exception of the wrapped function
-* `metric` - A dict containing the data to be send to the backend
-* `func_args` - Original args received by the wrapped function.
-* `func_kwargs` - Original kwargs received by the wrapped function.
+A function which is called after calling the decorated function and before sending the metrics. It receives the following arguments:
+
+* `response`: the returned value of the wrapped function
+* `exception`: exception raised the wrapped function, if any
+* `metric`: a dictionary containing the data to be sent to the backend
+* `func`: the decorated function itself
+* `func_args`: original args received by the wrapped function.
+* `func_kwargs`: original kwargs received by the wrapped function.
 
 From within a hook you can change the `name` if you want the metrics to be
 split into multiple series.
@@ -244,6 +247,23 @@ def celery_task(self, **kwargs):
     return True
 ```
 
+### Generator hook
+
+A generator hook works similarly to a simple hook, but also allows to run an arbitrary code just before the decorated function is called.
+It receives `func`, `func_args`, and `func_kwargs` as its arguments, whilst `response`, `exception`, and `metric` are sent later at the `yield` site:
+
+```python
+def generator_hook(func, func_args, func_kwargs) -> Generator[
+    None,
+    Tuple[Any, Optional[BaseException], Dict[str, Any]],  # response, exception, and metrics
+    Optional[Dict[str, Any]],
+]:
+    print("This runs before the decorated function")
+    (response, exception, metrics) = yield
+    print("This runs after the decorated function")
+    ...
+```
+
 ## Manually sending metrics
 
 You can also send any metric you have manually to the backend. These

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,40 @@ def hook(response, exception, metric, func, func_args, func_kwargs):
 
         with settings(hooks=[hook]):
             go(param1=param)
+
+    def test_generator_hook(self) -> None:
+        is_started = False
+
+        def generator_hook(func, func_args, func_kwargs) -> GeneratorHookReturnType:
+            assert func_args == (42,)
+            assert func_kwargs == {"bar": 100500}
+            assert not is_started, "the decorated function should not run just yet"
+            (response, _exception, _metrics) = yield
+            assert is_started
+            assert response == "response"
+            return {"key": "value"}
+
+        @time_execution(disable_default_hooks=True, extra_hooks=(generator_hook,))
+        def go(foo: int, *, bar: int) -> str:
+            nonlocal is_started
+            is_started = True
+            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]:
     """
@@ -90,4 +90,26 @@ def __call__(
         func_args: Tuple[Any, ...],
         func_kwargs: Dict[str, Any],
     ) -> Optional[Dict[str, Any]]:
-        ...
+        ...  # fmt:skip
+
+
+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:
+        ...  # fmt:skip

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,34 @@ 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 +80,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 +91,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