diff --git a/README.md b/README.md index 14772c8..854894b 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index ebd2b1f..61eab8e 100644 --- a/pyproject.toml +++ b/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 diff --git a/tests/test_hooks.py b/tests/test_hooks.py index d26e417..417c2aa 100644 --- a/tests/test_hooks.py +++ b/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() diff --git a/time_execution/decorator.py b/time_execution/decorator.py index 857b718..52ad3a2 100755 --- a/time_execution/decorator.py +++ b/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 diff --git a/time_execution/timed.py b/time_execution/timed.py index e7a5092..dc9dfa0 100644 --- a/time_execution/timed.py +++ b/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