BREAK: get rid of fqn_decorators.Decorator and improve the type hin…

…ting

Makefile

@@ -20,8 +20,7 @@ clean: pyclean
 venv: PYTHON?=python3.10
 venv:
 	$(PYTHON) -m venv venv
-	# FIXME: unpin when https://github.com/pypa/pip/issues/9215 is fixed
-	venv/bin/pip install -U "pip==20.2" -q
+	venv/bin/pip install -U pip -q
 	venv/bin/pip install -r requirements.txt '.[all]'
 
 ## Code style

setup.py

@@ -21,6 +21,7 @@
     install_requires=[
         "pkgsettings>=0.9.2",
         "fqn-decorators>=1.2.5,<3.0.0",
+        "typing-extensions>=4.5.0,<5.0.0",
     ],
     extras_require={
         "all": ["elasticsearch>=8.0.0,<9.0.0"],

tests/test_decorator_async.py

@@ -9,7 +9,7 @@
 @pytest.fixture
 def patch_backend(monkeypatch):
     m = Mock()
-    monkeypatch.setattr("time_execution.decorator.write_metric", m)
+    monkeypatch.setattr("time_execution.timed.write_metric", m)
     return m
 
 

tests/test_threaded_backend.py

@@ -11,10 +11,9 @@
 
 from tests.conftest import go
 from tests.test_base_backend import TestBaseBackend
-from time_execution import settings
+from time_execution import SHORT_HOSTNAME, settings
 from time_execution.backends import elasticsearch
 from time_execution.backends.threaded import ThreadedBackend
-from time_execution.decorator import SHORT_HOSTNAME
 
 from .test_elasticsearch import ELASTICSEARCH_URI, ElasticTestMixin
 

time_execution/__init__.py

@@ -1 +1,2 @@
-from .decorator import *  # noqa
+from .decorator import *  # noqa: F401, F403
+from .timed import SHORT_HOSTNAME  # noqa: F401

time_execution/decorator.py

@@ -1,91 +1,76 @@
 """
 Time Execution decorator
 """
-import socket
-import time
+from asyncio import iscoroutinefunction
+from functools import wraps
+from typing import Any, Callable, List, Optional, TypeVar, cast
 
-from fqn_decorators import Decorator
-from fqn_decorators.asynchronous import AsyncDecorator
+import fqn_decorators
 from pkgsettings import Settings
+from typing_extensions import overload
 
-SHORT_HOSTNAME = socket.gethostname()
+_F = TypeVar("_F", bound=Callable[..., Any])
 
 settings = Settings()
 settings.configure(backends=[], hooks=[], duration_field="value")
 
 
-def write_metric(name, **metric):
+def write_metric(name: str, **metric: Any) -> None:
     for backend in settings.backends:
         backend.write(name, **metric)
 
 
-def _apply_hooks(hooks, response, exception, metric, func, func_args, func_kwargs):
-    metadata = dict()
-    for hook in hooks:
-        hook_result = hook(
-            response=response,
-            exception=exception,
-            metric=metric,
-            func=func,
-            func_args=func_args,
-            func_kwargs=func_kwargs,
-        )
-
-        if hook_result:
-            metadata.update(hook_result)
-    return metadata
-
-
-class time_execution(Decorator):
-    def __init__(self, func=None, **params):
-        self.start_time = None
-        super(time_execution, self).__init__(func, **params)
-
-    def before(self):
-        self.start_time = time.time()
-
-    def after(self):
-        duration = round(time.time() - self.start_time, 3) * 1000
-
-        metric = {"name": self.fqn, settings.duration_field: duration, "hostname": SHORT_HOSTNAME}
-
-        origin = getattr(settings, "origin", None)
-        if origin:
-            metric["origin"] = origin
-
-        hooks = self.params.get("extra_hooks", [])
-        disable_default_hooks = self.params.get("disable_default_hooks", False)
-
-        if not 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 = _apply_hooks(
-            hooks=hooks,
-            response=self.result,
-            exception=self.get_exception(),
-            metric=metric,
-            func=self.func,
-            func_args=self.args,
-            func_kwargs=self.kwargs,
-        )
-
-        metric.update(metadata)
-        write_metric(**metric)
-
-    def get_exception(self):
-        """Retrieve the exception"""
-        if self.exc_info is None:
-            return
-
-        exc_type, exc_value, exc_tb = self.exc_info
-        if exc_value is None:
-            exc_value = exc_type()
-        if exc_value.__traceback__ is not exc_tb:
-            return exc_value.with_traceback(exc_tb)
-        return exc_value
-
-
-class time_execution_async(AsyncDecorator, time_execution):
-    pass
+@overload
+def time_execution(__wrapped: _F) -> _F:
+    """First-order (non-parametrized) decorator with the default FQN getter and hooks by default."""
+
+
+@overload
+def time_execution(
+    *,
+    get_fqn: Callable[[Any], str] = fqn_decorators.get_fqn,
+    extra_hooks: Optional[List] = None,
+    disable_default_hooks: bool = False,
+) -> Callable[[_F], _F]:
+    """
+    Second-order (parametrized) decorator.
+
+    Args:
+        get_fqn: custom FQN getter (uses `fqn-decorators` by default)
+        extra_hooks: additional hooks (next to defined in the settings)
+        disable_default_hooks: if `True`, disable the hooks set by the settings
+    """
+
+
+def time_execution(__wrapped=None, get_fqn: Callable[[Any], str] = fqn_decorators.get_fqn, **kwargs):
+    from time_execution.timed import Timed  # work around the circular dependency
+
+    def wrap(__wrapped: _F) -> _F:
+        fqn = get_fqn(__wrapped)
+
+        if not iscoroutinefunction(__wrapped):
+
+            @wraps(__wrapped)
+            def wrapper(*call_args, **call_kwargs):
+                with Timed(wrapped=__wrapped, call_args=call_args, call_kwargs=call_kwargs, fqn=fqn, **kwargs) as timed:
+                    timed.result = __wrapped(*call_args, **call_kwargs)
+                    return timed.result
+
+        else:
+
+            @wraps(__wrapped)
+            async def wrapper(*call_args, **call_kwargs):
+                with Timed(wrapped=__wrapped, call_args=call_args, call_kwargs=call_kwargs, fqn=fqn, **kwargs) as timed:
+                    timed.result = await __wrapped(*call_args, **call_kwargs)
+                    return timed.result
+
+        # Backwards compatibility with `Decorator`.
+        wrapper.fqn = fqn  # type: ignore[attr-defined]
+        wrapper.get_fqn = lambda: wrapper.fqn  # type: ignore[attr-defined]
+        return cast(_F, wrapper)
+
+    return wrap(__wrapped) if __wrapped is not None else wrap
+
+
+# `time_execution` supports async out of the box.
+time_execution_async = time_execution

time_execution/timed.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from contextlib import AbstractContextManager
+from socket import gethostname
+from timeit import default_timer
+from types import TracebackType
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+
+from time_execution import settings, write_metric
+
+SHORT_HOSTNAME = gethostname()
+
+
+class Timed(AbstractContextManager):
+    """
+    Both the sync and async decorators require the same logic around the wrapped function.
+    This context manager encapsulates the shared behaviour to avoid duplicating the code.
+    """
+
+    __slots__ = (
+        "result",
+        "_wrapped",
+        "_fqn",
+        "_extra_hooks",
+        "_disable_default_hooks",
+        "_call_args",
+        "_call_kwargs",
+        "_start_time",
+    )
+
+    def __init__(
+        self,
+        *,
+        wrapped: Callable[..., Any],
+        fqn: str,
+        call_args: Tuple[Any, ...],
+        call_kwargs: Dict[str, Any],
+        extra_hooks: Optional[List] = 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
+
+    def __enter__(self) -> Timed:
+        self._start_time = default_timer()
+        return self
+
+    def __exit__(
+        self,
+        __exc_type: Optional[Type[BaseException]],
+        __exc_val: Optional[BaseException],
+        __exc_tb: Optional[TracebackType],
+    ) -> None:
+        duration_millis = round(default_timer() - self._start_time, 3) * 1000.0
+
+        metric = {settings.duration_field: duration_millis, "hostname": SHORT_HOSTNAME}
+
+        origin = getattr(settings, "origin", None)
+        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)
+
+        metric.update(metadata)
+        write_metric(name=self._fqn, **metric)
+
+    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,
+            )
+            if hook_result:
+                metadata.update(hook_result)
+        return metadata