haiway 0.1.0__py3-none-any.whl

haiway/context/metrics.py

@@ -0,0 +1,329 @@
+from asyncio import Future, gather, get_event_loop
+from collections.abc import Callable
+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 haiway.state import Structure
+from haiway.utils import freeze
+
+__all__ = [
+    "ScopeMetrics",
+    "MetricsContext",
+]
+
+
+@final
+class ScopeMetrics:
+    def __init__(
+        self,
+        *,
+        trace_id: str | None,
+        scope: str,
+        logger: Logger | None,
+    ) -> None:
+        self.trace_id: str = trace_id or uuid4().hex
+        self._label: str = f"{self.trace_id}|{scope}" if scope else self.trace_id
+        self._logger: Logger = logger or getLogger(name=scope)
+        self._metrics: dict[type[Structure], Structure] = {}
+        self._nested: list[ScopeMetrics] = []
+        self._timestamp: float = monotonic()
+        self._completed: Future[float] = get_event_loop().create_future()
+
+        freeze(self)
+
+    def __del__(self) -> None:
+        self._complete()  # ensure completion on deinit
+
+    def __str__(self) -> str:
+        return self._label
+
+    def metrics(
+        self,
+        *,
+        merge: Callable[[Structure, Structure], Structure] = lambda lhs, rhs: lhs,
+    ) -> list[Structure]:
+        metrics: dict[type[Structure], Structure] = copy(self._metrics)
+        for metric in chain.from_iterable(nested.metrics(merge=merge) for nested in self._nested):
+            metric_type: type[Structure] = type(metric)
+            if current := metrics.get(metric_type):
+                metrics[metric_type] = merge(current, metric)
+
+            else:
+                metrics[metric_type] = metric
+
+        return list(metrics.values())
+
+    @overload
+    def read[Metric: Structure](
+        self,
+        metric: type[Metric],
+        /,
+    ) -> Metric | None: ...
+
+    @overload
+    def read[Metric: Structure](
+        self,
+        metric: type[Metric],
+        /,
+        default: Metric,
+    ) -> Metric: ...
+
+    def read[Metric: Structure](
+        self,
+        metric: type[Metric],
+        /,
+        default: Metric | None = None,
+    ) -> Metric | None:
+        return cast(Metric | None, self._metrics.get(metric, default))
+
+    def record[Metric: Structure](
+        self,
+        metric: Metric,
+        /,
+        *,
+        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 completed(self) -> bool:
+        return self._completed.done() and all(nested.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 _complete(self) -> None:
+        if self._completed.done():
+            return  # already completed
+
+        self._completed.set_result(monotonic() - self._timestamp)
+
+    def scope(
+        self,
+        name: str,
+        /,
+    ) -> Self:
+        nested: Self = self.__class__(
+            scope=name,
+            logger=self._logger,
+            trace_id=self.trace_id,
+        )
+        self._nested.append(nested)
+        return nested
+
+    def log(
+        self,
+        level: int,
+        message: str,
+        /,
+        *args: Any,
+        exception: BaseException | None = None,
+    ) -> None:
+        self._logger.log(
+            level,
+            f"[{self}] {message}",
+            *args,
+            exc_info=exception,
+        )
+
+
+@final
+class MetricsContext:
+    _context = ContextVar[ScopeMetrics]("MetricsContext")
+
+    @classmethod
+    def scope(
+        cls,
+        name: str,
+        /,
+        *,
+        trace_id: str | None = None,
+        logger: Logger | None = None,
+    ) -> Self:
+        try:
+            context: ScopeMetrics = cls._context.get()
+            if trace_id is None or context.trace_id == trace_id:
+                return cls(context.scope(name))
+
+            else:
+                return cls(
+                    ScopeMetrics(
+                        trace_id=trace_id,
+                        scope=name,
+                        logger=logger or context._logger,  # pyright: ignore[reportPrivateUsage]
+                    )
+                )
+        except LookupError:  # create metrics scope when missing yet
+            return cls(
+                ScopeMetrics(
+                    trace_id=trace_id,
+                    scope=name,
+                    logger=logger,
+                )
+            )
+
+    @classmethod
+    def record[Metric: Structure](
+        cls,
+        metric: Metric,
+        /,
+        *,
+        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)
+
+        except Exception as exc:
+            cls.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,
+    ) -> None:
+        self._metrics: ScopeMetrics = metrics
+        self._token: Token[ScopeMetrics] | None = None
+        self._started: float | None = None
+        self._finished: float | None = None
+
+    def __enter__(self) -> None:
+        assert (  # nosec: B101
+            self._token is None and self._started is None
+        ), "MetricsContext reentrance is not allowed"
+        self._token = MetricsContext._context.set(self._metrics)
+        self._started = monotonic()
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert (  # nosec: B101
+            self._token is not None and self._started is not None and self._finished is None
+        ), "Unbalanced MetricsContext context enter/exit"
+        self._finished = monotonic()
+        MetricsContext._context.reset(self._token)
+        self._token = None

haiway/context/state.py

@@ -0,0 +1,115 @@
+from collections.abc import Iterable
+from contextvars import ContextVar, Token
+from types import TracebackType
+from typing import Self, cast, final
+
+from haiway.context.types import MissingContext, MissingState
+from haiway.state import Structure
+from haiway.utils import freeze
+
+__all__ = [
+    "ScopeState",
+    "StateContext",
+]
+
+
+@final
+class ScopeState:
+    def __init__(
+        self,
+        state: Iterable[Structure],
+    ) -> None:
+        self._state: dict[type[Structure], Structure] = {
+            type(element): element for element in state
+        }
+        freeze(self)
+
+    def state[State: Structure](
+        self,
+        state: type[State],
+        /,
+        default: State | None = None,
+    ) -> State:
+        if state in self._state:
+            return cast(State, self._state[state])
+
+        elif default is not None:
+            return default
+
+        else:
+            try:
+                initialized: State = state()
+                self._state[state] = initialized
+                return initialized
+
+            except Exception as exc:
+                raise MissingState(
+                    f"{state.__qualname__} is not defined in current scope"
+                    " and failed to provide a default value"
+                ) from exc
+
+    def updated(
+        self,
+        state: Iterable[Structure],
+    ) -> Self:
+        if state:
+            return self.__class__(
+                [
+                    *self._state.values(),
+                    *state,
+                ]
+            )
+
+        else:
+            return self
+
+
+@final
+class StateContext:
+    _context = ContextVar[ScopeState]("StateContext")
+
+    @classmethod
+    def current[State: Structure](
+        cls,
+        state: type[State],
+        /,
+        default: State | None = None,
+    ) -> State:
+        try:
+            return cls._context.get().state(state, default=default)
+
+        except LookupError as exc:
+            raise MissingContext("StateContext requested but not defined!") from exc
+
+    @classmethod
+    def updated(
+        cls,
+        state: Iterable[Structure],
+        /,
+    ) -> Self:
+        try:
+            return cls(state=cls._context.get().updated(state=state))
+
+        except LookupError:  # create new context as a fallback
+            return cls(state=ScopeState(state))
+
+    def __init__(
+        self,
+        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)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert self._token is not None, "Unbalanced StateContext context exit"  # nosec: B101
+        StateContext._context.reset(self._token)
+        self._token = None

haiway/context/tasks.py

@@ -0,0 +1,65 @@
+from asyncio import Task, TaskGroup, get_event_loop
+from collections.abc import Callable, Coroutine
+from contextvars import ContextVar, Token, copy_context
+from types import TracebackType
+from typing import final
+
+__all__ = [
+    "TaskGroupContext",
+]
+
+
+@final
+class TaskGroupContext:
+    _context = ContextVar[TaskGroup]("TaskGroupContext")
+
+    @classmethod
+    def run[Result, **Arguments](
+        cls,
+        function: Callable[Arguments, Coroutine[None, None, Result]],
+        /,
+        *args: Arguments.args,
+        **kwargs: Arguments.kwargs,
+    ) -> Task[Result]:
+        try:
+            return cls._context.get().create_task(
+                function(*args, **kwargs),
+                context=copy_context(),
+            )
+
+        except LookupError:  # spawn task out of group as a fallback
+            return get_event_loop().create_task(
+                function(*args, **kwargs),
+                context=copy_context(),
+            )
+
+    def __init__(
+        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
+        await self._group.__aenter__()
+        self._token = TaskGroupContext._context.set(self._group)
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert self._token is not None, "Unbalanced TaskGroupContext context exit"  # nosec: B101
+        TaskGroupContext._context.reset(self._token)
+        self._token = None
+
+        try:
+            await self._group.__aexit__(
+                et=exc_type,
+                exc=exc_val,
+                tb=exc_tb,
+            )
+
+        except BaseException:
+            pass  # silence TaskGroup exceptions, if there was exception already we will get it

haiway/context/types.py

@@ -0,0 +1,17 @@
+__all__ = [
+    "MissingContext",
+    "MissingDependency",
+    "MissingState",
+]
+
+
+class MissingContext(Exception):
+    pass
+
+
+class MissingDependency(Exception):
+    pass
+
+
+class MissingState(Exception):
+    pass

haiway/helpers/__init__.py

@@ -0,0 +1,13 @@
+from haiway.helpers.asynchronous import asynchronous
+from haiway.helpers.cache import cached
+from haiway.helpers.retry import auto_retry
+from haiway.helpers.throttling import throttle
+from haiway.helpers.timeout import with_timeout
+
+__all__ = [
+    "asynchronous",
+    "auto_retry",
+    "cached",
+    "throttle",
+    "with_timeout",
+]