haiway 0.10.14__py3-none-any.whl → 0.10.16__py3-none-any.whl

haiway/context/identifier.py

@@ -0,0 +1,92 @@
+from contextvars import ContextVar, Token
+from types import TracebackType
+from typing import Any, Self, final
+from uuid import uuid4
+
+__all__ = [
+    "ScopeIdentifier",
+]
+
+
+@final
+class ScopeIdentifier:
+    _context = ContextVar[Self]("ScopeIdentifier")
+
+    @classmethod
+    def current_trace_id(cls) -> str:
+        try:
+            return ScopeIdentifier._context.get().trace_id
+
+        except LookupError as exc:
+            raise RuntimeError("Attempting to access scope identifier outside of scope") from exc
+
+    @classmethod
+    def scope(
+        cls,
+        label: str,
+        /,
+    ) -> Self:
+        current: Self
+        try:  # check for current scope
+            current = cls._context.get()
+
+        except LookupError:
+            # create root scope when missing
+            trace_id: str = uuid4().hex
+            return cls(
+                label=label,
+                scope_id=uuid4().hex,
+                parent_id=trace_id,  # trace_id is parent_id for root
+                trace_id=trace_id,
+            )
+
+        # create nested scope otherwise
+        return cls(
+            label=label,
+            scope_id=uuid4().hex,
+            parent_id=current.scope_id,
+            trace_id=current.trace_id,
+        )
+
+    def __init__(
+        self,
+        trace_id: str,
+        parent_id: str,
+        scope_id: str,
+        label: str,
+    ) -> None:
+        self.trace_id: str = trace_id
+        self.parent_id: str = parent_id
+        self.scope_id: str = scope_id
+        self.label: str = label
+        self.unique_name: str = f"[{trace_id}] [{label}] [{scope_id}]"
+
+    @property
+    def is_root(self) -> bool:
+        return self.trace_id == self.parent_id
+
+    def __str__(self) -> str:
+        return self.unique_name
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+
+        return self.scope_id == other.scope_id and self.trace_id == other.trace_id
+
+    def __hash__(self) -> int:
+        return hash(self.scope_id)
+
+    def __enter__(self) -> None:
+        assert not hasattr(self, "_token"), "Context reentrance is not allowed"  # nosec: B101
+        self._token: Token[ScopeIdentifier] = ScopeIdentifier._context.set(self)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
+        ScopeIdentifier._context.reset(self._token)
+        del self._token

haiway/context/logging.py

@@ -0,0 +1,176 @@
+from contextvars import ContextVar, Token
+from logging import DEBUG, ERROR, INFO, WARNING, Logger, getLogger
+from time import monotonic
+from types import TracebackType
+from typing import Any, Self, final
+
+from haiway.context.identifier import ScopeIdentifier
+
+__all__ = [
+    "LoggerContext",
+]
+
+
+@final
+class LoggerContext:
+    _context = ContextVar[Self]("LoggerContext")
+
+    @classmethod
+    def scope(
+        cls,
+        scope: ScopeIdentifier,
+        /,
+        *,
+        logger: Logger | None,
+    ) -> Self:
+        current: Self
+        try:  # check for current scope
+            current = cls._context.get()
+
+        except LookupError:
+            # create root scope when missing
+            return cls(
+                scope=scope,
+                logger=logger,
+            )
+
+        # create nested scope otherwise
+        return cls(
+            scope=scope,
+            logger=logger or current._logger,
+        )
+
+    @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,
+        scope: ScopeIdentifier,
+        logger: Logger | None,
+    ) -> None:
+        self._prefix: str = scope.unique_name
+        self._logger: Logger = logger or getLogger(name=scope.label)
+
+    def log(
+        self,
+        level: int,
+        message: str,
+        /,
+        *args: Any,
+        exception: BaseException | None = None,
+    ) -> None:
+        self._logger.log(
+            level,
+            f"{self._prefix} {message}",
+            *args,
+            exc_info=exception,
+        )
+
+    def __enter__(self) -> None:
+        assert not hasattr(self, "_token"), "Context reentrance is not allowed"  # nosec: B101
+        assert not hasattr(self, "_entered"), "Context reentrance is not allowed"  # nosec: B101
+        self._entered: float = monotonic()
+        self._token: Token[LoggerContext] = LoggerContext._context.set(self)
+        self.log(DEBUG, "Entering context...")
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
+        LoggerContext._context.reset(self._token)
+        del self._token
+        self.log(DEBUG, f"...exiting context after {monotonic() - self._entered:.2f}s")
+        del self._entered

haiway/context/metrics.py

@@ -0,0 +1,165 @@
+from contextvars import ContextVar, Token
+from types import TracebackType
+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
+
+__all__ = [
+    "MetricsContext",
+    "MetricsHandler",
+    "MetricsReading",
+    "MetricsRecording",
+    "MetricsScopeEntering",
+    "MetricsScopeExiting",
+]
+
+
+@runtime_checkable
+class MetricsRecording(Protocol):
+    def __call__(
+        self,
+        scope: ScopeIdentifier,
+        /,
+        metric: State,
+    ) -> None: ...
+
+
+@runtime_checkable
+class MetricsReading(Protocol):
+    async def __call__[Metric: State](
+        self,
+        scope: ScopeIdentifier,
+        /,
+        *,
+        metric: type[Metric],
+        merged: bool,
+    ) -> Metric | None: ...
+
+
+@runtime_checkable
+class MetricsScopeEntering(Protocol):
+    def __call__[Metric: State](
+        self,
+        scope: ScopeIdentifier,
+        /,
+    ) -> None: ...
+
+
+@runtime_checkable
+class MetricsScopeExiting(Protocol):
+    def __call__[Metric: State](
+        self,
+        scope: ScopeIdentifier,
+        /,
+    ) -> None: ...
+
+
+class MetricsHandler(State):
+    record: MetricsRecording
+    read: MetricsReading
+    enter_scope: MetricsScopeEntering
+    exit_scope: MetricsScopeExiting
+
+
+@final
+class MetricsContext:
+    _context = ContextVar[Self]("MetricsContext")
+
+    @classmethod
+    def scope(
+        cls,
+        scope: ScopeIdentifier,
+        /,
+        *,
+        metrics: MetricsHandler | None,
+    ) -> Self:
+        current: Self
+        try:  # check for current scope
+            current = cls._context.get()
+
+        except LookupError:
+            # create root scope when missing
+            return cls(
+                scope=scope,
+                metrics=metrics,
+            )
+
+        # create nested scope otherwise
+        return cls(
+            scope=scope,
+            metrics=metrics or current._metrics,
+        )
+
+    @classmethod
+    def record(
+        cls,
+        metric: State,
+        /,
+    ) -> None:
+        try:  # catch exceptions - we don't wan't to blow up on metrics
+            metrics: Self = cls._context.get()
+
+            if metrics._metrics is not None:
+                metrics._metrics.record(
+                    metrics._scope,
+                    metric,
+                )
+
+        except Exception as exc:
+            LoggerContext.log_error(
+                "Failed to record metric: %s",
+                type(metric).__qualname__,
+                exception=exc,
+            )
+
+    @classmethod
+    async def read[Metric: State](
+        cls,
+        metric: type[Metric],
+        /,
+        merged: bool,
+    ) -> Metric | None:
+        try:  # catch exceptions - we don't wan't to blow up on metrics
+            metrics: Self = cls._context.get()
+
+            if metrics._metrics is not None:
+                return await metrics._metrics.read(
+                    metrics._scope,
+                    metric=metric,
+                    merged=merged,
+                )
+
+        except Exception as exc:
+            LoggerContext.log_error(
+                "Failed to read metric: %s",
+                metric.__qualname__,
+                exception=exc,
+            )
+
+    def __init__(
+        self,
+        scope: ScopeIdentifier,
+        metrics: MetricsHandler | None,
+    ) -> None:
+        self._scope: ScopeIdentifier = scope
+        self._metrics: MetricsHandler | None = metrics
+
+    def __enter__(self) -> None:
+        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,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
+        MetricsContext._context.reset(self._token)
+        del self._token
+        if self._metrics is not None:
+            self._metrics.exit_scope(self._scope)

haiway/context/state.py

@@ -0,0 +1,113 @@
+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 State
+from haiway.utils import freeze
+
+__all__ = [
+    "ScopeState",
+    "StateContext",
+]
+
+
+@final
+class ScopeState:
+    def __init__(
+        self,
+        state: Iterable[State],
+    ) -> None:
+        self._state: dict[type[State], State] = {type(element): element for element in state}
+        freeze(self)
+
+    def state[StateType: State](
+        self,
+        state: type[StateType],
+        /,
+        default: StateType | None = None,
+    ) -> StateType:
+        if state in self._state:
+            return cast(StateType, self._state[state])
+
+        elif default is not None:
+            return default
+
+        else:
+            try:
+                initialized: StateType = 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[State],
+    ) -> Self:
+        if state:
+            return self.__class__(
+                [
+                    *self._state.values(),
+                    *state,
+                ]
+            )
+
+        else:
+            return self
+
+
+@final
+class StateContext:
+    _context = ContextVar[ScopeState]("StateContext")
+
+    @classmethod
+    def current[StateType: State](
+        cls,
+        state: type[StateType],
+        /,
+        default: StateType | None = None,
+    ) -> StateType:
+        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[State],
+        /,
+    ) -> Self:
+        try:
+            # update current scope context
+            return cls(state=cls._context.get().updated(state=state))
+
+        except LookupError:  # create root scope when missing
+            return cls(state=ScopeState(state))
+
+    def __init__(
+        self,
+        state: ScopeState,
+    ) -> None:
+        self._state: ScopeState = state
+
+    def __enter__(self) -> None:
+        assert not hasattr(self, "_token"), "Context reentrance is not allowed"  # nosec: B101
+        self._token: Token[ScopeState] = StateContext._context.set(self._state)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        assert hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
+        StateContext._context.reset(self._token)
+        del self._token

haiway/context/tasks.py

@@ -0,0 +1,64 @@
+from asyncio import Task, TaskGroup, get_event_loop
+from collections.abc import Callable, Coroutine
+from contextvars import ContextVar, 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()
+
+    async def __aenter__(self) -> None:
+        assert not hasattr(self, "_token"), "Context 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 hasattr(self, "_token"), "Unbalanced context enter/exit"  # nosec: B101
+        TaskGroupContext._context.reset(self._token)
+        del self._token
+
+        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,12 @@
+__all__ = [
+    "MissingContext",
+    "MissingState",
+]
+
+
+class MissingContext(Exception):
+    pass
+
+
+class MissingState(Exception):
+    pass

haiway/helpers/__init__.py

@@ -0,0 +1,21 @@
+from haiway.helpers.asynchrony import asynchronous, wrap_async
+from haiway.helpers.caching import cache
+from haiway.helpers.metrics import MetricsHolder, MetricsLogger
+from haiway.helpers.retries import retry
+from haiway.helpers.throttling import throttle
+from haiway.helpers.timeouted import timeout
+from haiway.helpers.tracing import ArgumentsTrace, ResultTrace, traced
+
+__all__ = [
+    "ArgumentsTrace",
+    "MetricsHolder",
+    "MetricsLogger",
+    "ResultTrace",
+    "asynchronous",
+    "cache",
+    "retry",
+    "throttle",
+    "timeout",
+    "traced",
+    "wrap_async",
+]