NEW: generator hooks ✨

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