From c7119fbf108f0cfa3f330c9cee64f0f00953fb98 Mon Sep 17 00:00:00 2001 From: Pavel Perestoronin Date: Thu, 25 Jan 2024 16:22:06 +0100 Subject: [PATCH] =?UTF-8?q?NEW:=20generator=20hooks=20=E2=9C=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### 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 | 32 +++++++++++++++---- pyproject.toml | 3 ++ tests/test_hooks.py | 39 ++++++++++++++++++++++- time_execution/decorator.py | 28 +++++++++++++++-- time_execution/timed.py | 63 +++++++++++++++++++++++-------------- 5 files changed, 132 insertions(+), 33 deletions(-) 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..344bcaf 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]: """ @@ -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: + ... diff --git a/time_execution/timed.py b/time_execution/timed.py index e7a5092..4c4683a 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,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