Skip to content

Commit

Permalink
NEW: better typing for hooks πŸ§‘β€πŸ’»
Browse files Browse the repository at this point in the history
  • Loading branch information
Pavel Perestoronin authored and eigenein committed Jan 25, 2024
1 parent f921346 commit 18088ef
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 12 deletions.
31 changes: 24 additions & 7 deletions time_execution/decorator.py
Original file line number Diff line number Diff line change
@@ -1,18 +1,20 @@
"""
Time Execution decorator
"""
"""Time Execution decorator"""

from __future__ import annotations

from asyncio import iscoroutinefunction
from collections.abc import Iterable
from functools import wraps
from typing import Any, Callable, List, Optional, TypeVar, cast
from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, cast

import fqn_decorators
from pkgsettings import Settings
from typing_extensions import overload
from typing_extensions import Protocol, overload

_F = TypeVar("_F", bound=Callable[..., Any])

settings = Settings()
settings.configure(backends=[], hooks=[], duration_field="value")
settings.configure(backends=(), hooks=(), duration_field="value")


def write_metric(name: str, **metric: Any) -> None:
Expand All @@ -29,7 +31,7 @@ def time_execution(__wrapped: _F) -> _F:
def time_execution(
*,
get_fqn: Callable[[Any], str] = fqn_decorators.get_fqn,
extra_hooks: Optional[List] = None,
extra_hooks: Optional[Iterable[Hook]] = None,
disable_default_hooks: bool = False,
) -> Callable[[_F], _F]:
"""
Expand Down Expand Up @@ -74,3 +76,18 @@ async def wrapper(*call_args, **call_kwargs):

# `time_execution` supports async out of the box.
time_execution_async = time_execution


class Hook(Protocol):
"""Hook callback protocol."""

def __call__(
self,
response: Any,
exception: Optional[BaseException],
metric: Dict[str, Any],
func: Callable[..., Any],
func_args: Tuple[Any, ...],
func_kwargs: Dict[str, Any],
) -> Optional[Dict[str, Any]]:
...
11 changes: 6 additions & 5 deletions time_execution/timed.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
from __future__ import annotations

from collections.abc import Iterable
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 typing import Any, Callable, Dict, Optional, Tuple, Type

from time_execution import settings, write_metric
from time_execution import Hook, settings, write_metric

SHORT_HOSTNAME = gethostname()

Expand Down Expand Up @@ -35,7 +36,7 @@ def __init__(
fqn: str,
call_args: Tuple[Any, ...],
call_kwargs: Dict[str, Any],
extra_hooks: Optional[List] = None,
extra_hooks: Optional[Iterable[Hook]] = None,
disable_default_hooks: bool = False,
) -> None:
self.result: Optional[Any] = None
Expand Down Expand Up @@ -64,9 +65,9 @@ def __exit__(
if origin:
metric["origin"] = origin

hooks = self._extra_hooks or []
hooks = self._extra_hooks or ()
if not self._disable_default_hooks:
hooks = settings.hooks + hooks
hooks = (*settings.hooks, *hooks)

# Apply the registered hooks, and collect the metadata they might
# return to be stored with the metrics.
Expand Down

0 comments on commit 18088ef

Please sign in to comment.