haiway 0.7.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

haiway/context/metrics.py

@@ -1,384 +1,118 @@
-from asyncio import (
-    AbstractEventLoop,
-    Future,
-    gather,
-    get_event_loop,
-    iscoroutinefunction,
-    run_coroutine_threadsafe,
-)
-from collections.abc import Callable, Coroutine
 from contextvars import ContextVar, Token
-from copy import copy
-from itertools import chain
-from logging import DEBUG, ERROR, INFO, WARNING, Logger, getLogger
-from time import monotonic
 from types import TracebackType
-from typing import Any, Self, cast, final, overload
-from uuid import uuid4
+from typing import Protocol, Self, final, runtime_checkable
 
+from haiway.context.identifier import ScopeIdentifier
+from haiway.context.logging import LoggerContext
 from haiway.state import State
-from haiway.types import MISSING, Missing, not_missing
-from haiway.utils import freeze
 
 __all__ = [
     "MetricsContext",
-    "ScopeMetrics",
+    "MetricsHandler",
+    "MetricsRecording",
+    "MetricsScopeEntering",
+    "MetricsScopeExiting",
 ]
 
 
-@final
-class ScopeMetrics:
-    def __init__(
-        self,
-        *,
-        trace_id: str | None,
-        scope: str,
-        logger: Logger | None,
-        parent: Self | None,
-        completion: Callable[[Self], Coroutine[None, None, None]] | Callable[[Self], None] | None,
-    ) -> None:
-        self.trace_id: str = trace_id or uuid4().hex
-        self.identifier: str = uuid4().hex
-        self.label: str = scope
-        self._logger_prefix: str = (
-            f"[{self.trace_id}] [{scope}] [{self.identifier}]"
-            if scope
-            else f"[{self.trace_id}] [{self.identifier}]"
-        )
-        self._logger: Logger = logger or getLogger(name=scope)
-        self._parent: Self | None = parent if parent else None
-        self._metrics: dict[type[State], State] = {}
-        self._nested: list[ScopeMetrics] = []
-        self._timestamp: float = monotonic()
-        self._finished: bool = False
-        self._loop: AbstractEventLoop = get_event_loop()
-        self._completed: Future[float] = self._loop.create_future()
-
-        if parent := parent:
-            parent._nested.append(self)
-
-        freeze(self)
-
-        if completion := completion:
-            metrics: Self = self
-            if iscoroutinefunction(completion):
-
-                def callback(_: Future[float]) -> None:
-                    run_coroutine_threadsafe(
-                        completion(metrics),
-                        metrics._loop,
-                    )
-
-            else:
-
-                def callback(_: Future[float]) -> None:
-                    completion(metrics)
-
-            self._completed.add_done_callback(callback)
-
-    def __del__(self) -> None:
-        assert self.is_completed, "Deinitializing not completed scope metrics"  # nosec: B101
-
-    def __str__(self) -> str:
-        return f"{self.label}[{self.identifier}]@[{self.trace_id}]"
-
-    def metrics(
-        self,
-        *,
-        merge: Callable[[State | Missing, State], State | Missing] | None = None,
-    ) -> list[State]:
-        if not merge:
-            return list(self._metrics.values())
-
-        metrics: dict[type[State], State] = copy(self._metrics)
-        for metric in chain.from_iterable(nested.metrics(merge=merge) for nested in self._nested):
-            metric_type: type[State] = type(metric)
-            merged: State | Missing = merge(
-                metrics.get(  # current
-                    metric_type,
-                    MISSING,
-                ),
-                metric,  # received
-            )
-
-            if not_missing(merged):
-                metrics[metric_type] = merged
-
-        return list(metrics.values())
-
-    @overload
-    def read[Metric: State](
+@runtime_checkable
+class MetricsRecording(Protocol):
+    def __call__(
         self,
-        metric: type[Metric],
+        scope: ScopeIdentifier,
         /,
-    ) -> Metric | None: ...
+        metric: State,
+    ) -> None: ...
 
-    @overload
-    def read[Metric: State](
-        self,
-        metric: type[Metric],
-        /,
-        default: Metric,
-    ) -> Metric: ...
 
-    def read[Metric: State](
+@runtime_checkable
+class MetricsScopeEntering(Protocol):
+    def __call__[Metric: State](
         self,
-        metric: type[Metric],
+        scope: ScopeIdentifier,
         /,
-        default: Metric | None = None,
-    ) -> Metric | None:
-        return cast(Metric | None, self._metrics.get(metric, default))
+    ) -> None: ...
 
-    def record[Metric: State](
+
+@runtime_checkable
+class MetricsScopeExiting(Protocol):
+    def __call__[Metric: State](
         self,
-        metric: Metric,
+        scope: ScopeIdentifier,
         /,
-        *,
-        merge: Callable[[Metric, Metric], Metric] = lambda lhs, rhs: rhs,
-    ) -> None:
-        assert not self._completed.done(), "Can't record using completed metrics scope"  # nosec: B101
-        metric_type: type[Metric] = type(metric)
-        if current := self._metrics.get(metric_type):
-            self._metrics[metric_type] = merge(cast(Metric, current), metric)
-
-        else:
-            self._metrics[metric_type] = metric
-
-    @property
-    def is_completed(self) -> bool:
-        return self._completed.done() and all(nested.is_completed for nested in self._nested)
-
-    @property
-    def time(self) -> float:
-        if self._completed.done():
-            return self._completed.result()
-
-        else:
-            return monotonic() - self._timestamp
-
-    async def wait(self) -> None:
-        await gather(
-            self._completed,
-            *[nested.wait() for nested in self._nested],
-            return_exceptions=False,
-        )
-
-    def _finish(self) -> None:
-        assert (  # nosec: B101
-            not self._completed.done()
-        ), "Invalid state - called finish on already completed scope"
-
-        assert (  # nosec: B101
-            not self._finished
-        ), "Invalid state - called completion on already finished scope"
-
-        self._finished = True  # self is now finished
-
-        self._complete_if_able()
-
-    def _complete_if_able(self) -> None:
-        assert (  # nosec: B101
-            not self._completed.done()
-        ), "Invalid state - called complete on already completed scope"
+    ) -> None: ...
 
-        if not self._finished:
-            return  # wait for finishing self
 
-        if any(not nested.is_completed for nested in self._nested):
-            return  # wait for completing all nested scopes
-
-        # set completion time
-        self._completed.set_result(monotonic() - self._timestamp)
-
-        # notify parent about completion
-        if parent := self._parent:
-            parent._complete_if_able()
-
-    def log(
-        self,
-        level: int,
-        message: str,
-        /,
-        *args: Any,
-        exception: BaseException | None = None,
-    ) -> None:
-        self._logger.log(
-            level,
-            f"{self._logger_prefix} {message}",
-            *args,
-            exc_info=exception,
-        )
+class MetricsHandler(State):
+    record: MetricsRecording
+    enter_scope: MetricsScopeEntering
+    exit_scope: MetricsScopeExiting
 
 
 @final
 class MetricsContext:
-    _context = ContextVar[ScopeMetrics]("MetricsContext")
+    _context = ContextVar[Self]("MetricsContext")
 
     @classmethod
     def scope(
         cls,
-        name: str,
+        scope: ScopeIdentifier,
         /,
         *,
-        trace_id: str | None = None,
-        logger: Logger | None = None,
-        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]]
-        | Callable[[ScopeMetrics], None]
-        | None,
+        metrics: MetricsHandler | None,
     ) -> Self:
-        current: ScopeMetrics
-        try:  # check for current scope context
+        current: Self
+        try:  # check for current scope
             current = cls._context.get()
 
         except LookupError:
-            # create metrics scope when missing yet
+            # create root scope when missing
             return cls(
-                ScopeMetrics(
-                    trace_id=trace_id,
-                    scope=name,
-                    logger=logger,
-                    parent=None,
-                    completion=completion,
-                )
+                scope=scope,
+                metrics=metrics,
             )
 
-        # or create nested metrics otherwise
+        # create nested scope otherwise
         return cls(
-            ScopeMetrics(
-                trace_id=trace_id,
-                scope=name,
-                logger=logger or current._logger,  # pyright: ignore[reportPrivateUsage]
-                parent=current,
-                completion=completion,
-            )
+            scope=scope,
+            metrics=metrics or current._metrics,
         )
 
     @classmethod
-    def record[Metric: State](
+    def record(
         cls,
-        metric: Metric,
+        metric: State,
         /,
-        *,
-        merge: Callable[[Metric, Metric], Metric] = lambda lhs, rhs: rhs,
     ) -> None:
         try:  # catch exceptions - we don't wan't to blow up on metrics
-            cls._context.get().record(metric, merge=merge)
+            metrics: Self = cls._context.get()
+
+            if metrics._metrics is not None:
+                metrics._metrics.record(
+                    metrics._scope,
+                    metric,
+                )
 
         except Exception as exc:
-            cls.log_error(
+            LoggerContext.log_error(
                 "Failed to record metric: %s",
                 type(metric).__qualname__,
                 exception=exc,
             )
 
-    # - LOGS -
-
-    @classmethod
-    def log_error(
-        cls,
-        message: str,
-        /,
-        *args: Any,
-        exception: BaseException | None = None,
-    ) -> None:
-        try:
-            cls._context.get().log(
-                ERROR,
-                message,
-                *args,
-                exception=exception,
-            )
-
-        except LookupError:
-            getLogger().log(
-                ERROR,
-                message,
-                *args,
-                exc_info=exception,
-            )
-
-    @classmethod
-    def log_warning(
-        cls,
-        message: str,
-        /,
-        *args: Any,
-        exception: Exception | None = None,
-    ) -> None:
-        try:
-            cls._context.get().log(
-                WARNING,
-                message,
-                *args,
-                exception=exception,
-            )
-
-        except LookupError:
-            getLogger().log(
-                WARNING,
-                message,
-                *args,
-                exc_info=exception,
-            )
-
-    @classmethod
-    def log_info(
-        cls,
-        message: str,
-        /,
-        *args: Any,
-    ) -> None:
-        try:
-            cls._context.get().log(
-                INFO,
-                message,
-                *args,
-            )
-
-        except LookupError:
-            getLogger().log(
-                INFO,
-                message,
-                *args,
-            )
-
-    @classmethod
-    def log_debug(
-        cls,
-        message: str,
-        /,
-        *args: Any,
-        exception: Exception | None = None,
-    ) -> None:
-        try:
-            cls._context.get().log(
-                DEBUG,
-                message,
-                *args,
-                exception=exception,
-            )
-
-        except LookupError:
-            getLogger().log(
-                DEBUG,
-                message,
-                *args,
-                exc_info=exception,
-            )
-
     def __init__(
         self,
-        metrics: ScopeMetrics,
+        scope: ScopeIdentifier,
+        metrics: MetricsHandler | None,
     ) -> None:
-        self._metrics: ScopeMetrics = metrics
-        self._token: Token[ScopeMetrics] | None = None
+        self._scope: ScopeIdentifier = scope
+        self._metrics: MetricsHandler | None = metrics
 
     def __enter__(self) -> None:
-        assert (  # nosec: B101
-            self._token is None and not self._metrics._finished  # pyright: ignore[reportPrivateUsage]
-        ), "MetricsContext reentrance is not allowed"
-        self._token = MetricsContext._context.set(self._metrics)
-        self._metrics.log(INFO, "Started...")
+        assert not hasattr(self, "_token"), "Context reentrance is not allowed"  # nosec: B101
+        self._token: Token[MetricsContext] = MetricsContext._context.set(self)
+        if self._metrics is not None:
+            self._metrics.enter_scope(self._scope)
 
     def __exit__(
         self,
@@ -386,10 +120,8 @@ class MetricsContext:
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        assert (  # nosec: B101
-            self._token is not None
-        ), "Unbalanced MetricsContext context enter/exit"
+        assert hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
         MetricsContext._context.reset(self._token)
-        self._metrics._finish()  # pyright: ignore[reportPrivateUsage]
-        self._token = None
-        self._metrics.log(INFO, f"...finished after {self._metrics.time:.2f}s")
+        del self._token
+        if self._metrics is not None:
+            self._metrics.exit_scope(self._scope)

haiway/context/state.py

@@ -86,9 +86,10 @@ class StateContext:
         /,
     ) -> Self:
         try:
+            # update current scope context
             return cls(state=cls._context.get().updated(state=state))
 
-        except LookupError:  # create new context as a fallback
+        except LookupError:  # create root scope when missing
             return cls(state=ScopeState(state))
 
     def __init__(
@@ -96,11 +97,10 @@ class StateContext:
         state: ScopeState,
     ) -> None:
         self._state: ScopeState = state
-        self._token: Token[ScopeState] | None = None
 
     def __enter__(self) -> None:
-        assert self._token is None, "StateContext reentrance is not allowed"  # nosec: B101
-        self._token = StateContext._context.set(self._state)
+        assert not hasattr(self, "_token"), "Context reentrance is not allowed"  # nosec: B101
+        self._token: Token[ScopeState] = StateContext._context.set(self._state)
 
     def __exit__(
         self,
@@ -108,6 +108,6 @@ class StateContext:
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        assert self._token is not None, "Unbalanced StateContext context exit"  # nosec: B101
+        assert hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
         StateContext._context.reset(self._token)
-        self._token = None
+        del self._token

haiway/context/tasks.py

@@ -1,6 +1,6 @@
 from asyncio import Task, TaskGroup, get_event_loop
 from collections.abc import Callable, Coroutine
-from contextvars import ContextVar, Token, copy_context
+from contextvars import ContextVar, copy_context
 from types import TracebackType
 from typing import final
 
@@ -37,10 +37,9 @@ class TaskGroupContext:
         self,
     ) -> None:
         self._group: TaskGroup = TaskGroup()
-        self._token: Token[TaskGroup] | None = None
 
     async def __aenter__(self) -> None:
-        assert self._token is None, "TaskGroupContext reentrance is not allowed"  # nosec: B101
+        assert not hasattr(self, "_token"), "Context reentrance is not allowed"  # nosec: B101
         await self._group.__aenter__()
         self._token = TaskGroupContext._context.set(self._group)
 
@@ -50,9 +49,9 @@ class TaskGroupContext:
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        assert self._token is not None, "Unbalanced TaskGroupContext context exit"  # nosec: B101
+        assert hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
         TaskGroupContext._context.reset(self._token)
-        self._token = None
+        del self._token
 
         try:
             await self._group.__aexit__(

haiway/helpers/__init__.py

@@ -1,5 +1,6 @@
 from haiway.helpers.asynchrony import asynchronous, wrap_async
 from haiway.helpers.caching import cache
+from haiway.helpers.metrics import MetricsLogger
 from haiway.helpers.retries import retry
 from haiway.helpers.throttling import throttle
 from haiway.helpers.timeouted import timeout
@@ -7,6 +8,7 @@ from haiway.helpers.tracing import ArgumentsTrace, ResultTrace, traced
 
 __all__ = [
     "ArgumentsTrace",
+    "MetricsLogger",
     "ResultTrace",
     "asynchronous",
     "cache",