haiway 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

haiway/context/access.py

@@ -1,12 +1,16 @@
 from asyncio import (
+    CancelledError,
     Task,
     current_task,
 )
 from collections.abc import (
+    AsyncGenerator,
+    AsyncIterator,
     Callable,
     Coroutine,
     Iterable,
 )
+from contextvars import Context, copy_context
 from logging import Logger
 from types import TracebackType
 from typing import Any, final
@@ -32,32 +36,28 @@ class ScopeContext:
         logger: Logger | None,
         state: tuple[State, ...],
         disposables: Disposables | None,
-        task_group: TaskGroupContext,
-        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]] | None,
+        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]]
+        | Callable[[ScopeMetrics], None]
+        | None,
     ) -> None:
-        self._task_group: TaskGroupContext = task_group
-        self._logger: Logger | None = logger
-        self._trace_id: str | None = trace_id
-        self._name: str = name
+        self._task_group_context: TaskGroupContext = TaskGroupContext()
+        # postponing state creation to include disposables if needed
         self._state_context: StateContext
         self._state: tuple[State, ...] = state
         self._disposables: Disposables | None = disposables
-        self._metrics_context: MetricsContext
-        self._completion: Callable[[ScopeMetrics], Coroutine[None, None, None]] | None = completion
+        # pre-building metrics context to ensure nested context registering
+        self._metrics_context: MetricsContext = MetricsContext.scope(
+            name,
+            logger=logger,
+            trace_id=trace_id,
+            completion=completion,
+        )
 
         freeze(self)
 
     def __enter__(self) -> None:
-        assert self._completion is None, "Can't enter synchronous context with completion"  # nosec: B101
         assert self._disposables is None, "Can't enter synchronous context with disposables"  # nosec: B101
-
         self._state_context = StateContext.updated(self._state)
-        self._metrics_context = MetricsContext.scope(
-            self._name,
-            logger=self._logger,
-            trace_id=self._trace_id,
-        )
-
         self._state_context.__enter__()
         self._metrics_context.__enter__()
 
@@ -80,9 +80,9 @@ class ScopeContext:
         )
 
     async def __aenter__(self) -> None:
-        await self._task_group.__aenter__()
+        await self._task_group_context.__aenter__()
 
-        if self._disposables:
+        if self._disposables is not None:
             self._state_context = StateContext.updated(
                 (*self._state, *await self._disposables.__aenter__())
             )
@@ -90,12 +90,6 @@ class ScopeContext:
         else:
             self._state_context = StateContext.updated(self._state)
 
-        self._metrics_context = MetricsContext.scope(
-            self._name,
-            logger=self._logger,
-            trace_id=self._trace_id,
-        )
-
         self._state_context.__enter__()
         self._metrics_context.__enter__()
 
@@ -105,14 +99,14 @@ class ScopeContext:
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        if self._disposables:
+        if self._disposables is not None:
             await self._disposables.__aexit__(
                 exc_type=exc_type,
                 exc_val=exc_val,
                 exc_tb=exc_tb,
             )
 
-        await self._task_group.__aexit__(
+        await self._task_group_context.__aexit__(
             exc_type=exc_type,
             exc_val=exc_val,
             exc_tb=exc_tb,
@@ -130,9 +124,6 @@ class ScopeContext:
             exc_tb=exc_tb,
         )
 
-        if completion := self._completion:
-            await completion(self._metrics_context._metrics)  # pyright: ignore[reportPrivateUsage]
-
 
 @final
 class ctx:
@@ -144,7 +135,9 @@ class ctx:
         disposables: Disposables | Iterable[Disposable] | None = None,
         logger: Logger | None = None,
         trace_id: str | None = None,
-        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]] | None = None,
+        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]]
+        | Callable[[ScopeMetrics], None]
+        | None = None,
     ) -> ScopeContext:
         """
         Access scope context with given parameters. When called within an existing context\
@@ -173,7 +166,7 @@ class ctx:
              provided current identifier will be used if any, otherwise it random id will\
              be generated
 
-        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]] | None = None
+        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]] | Callable[[ScopeMetrics], None] | None = None
             completion callback called on exit from the scope granting access to finished\
              scope metrics. Completion is called outside of the context when its metrics is\
              already finished. Make sure to avoid any long operations within the completion.
@@ -182,7 +175,7 @@ class ctx:
         -------
         ScopeContext
             context object intended to enter context manager with it
-        """
+        """  # noqa: E501
 
         resolved_disposables: Disposables | None
         match disposables:
@@ -201,7 +194,6 @@ class ctx:
             logger=logger,
             state=state,
             disposables=resolved_disposables,
-            task_group=TaskGroupContext(),
             completion=completion,
         )
 
@@ -257,6 +249,62 @@ class ctx:
 
         return TaskGroupContext.run(function, *args, **kwargs)
 
+    @staticmethod
+    def stream[Result, **Arguments](
+        source: Callable[Arguments, AsyncGenerator[Result, None]],
+        /,
+        *args: Arguments.args,
+        **kwargs: Arguments.kwargs,
+    ) -> AsyncIterator[Result]:
+        """
+        Stream results produced by a generator within the proper context state.
+
+        Parameters
+        ----------
+        source: Callable[Arguments, AsyncGenerator[Result, None]]
+            generator streamed as the result
+
+        *args: Arguments.args
+            positional arguments passed to generator call
+
+        **kwargs: Arguments.kwargs
+            keyword arguments passed to generator call
+
+        Returns
+        -------
+        AsyncIterator[Result]
+            iterator for accessing generated results
+        """
+
+        # prepare context snapshot
+        context_snapshot: Context = copy_context()
+
+        # prepare nested context
+        streaming_context: ScopeContext = ctx.scope(
+            getattr(
+                source,
+                "__name__",
+                "streaming",
+            )
+        )
+
+        async def generator() -> AsyncGenerator[Result, None]:
+            async with streaming_context:
+                async for result in source(*args, **kwargs):
+                    yield result
+
+        # finally return it as an iterator
+        return context_snapshot.run(generator)
+
+    @staticmethod
+    def check_cancellation() -> None:
+        """
+        Check if current asyncio task is cancelled, raises CancelledError if so.
+        """
+
+        if (task := current_task()) and task.cancelled():
+            raise CancelledError()
+
     @staticmethod
     def cancel() -> None:
         """

haiway/context/metrics.py

@@ -1,5 +1,12 @@
-from asyncio import Future, gather, get_event_loop
-from collections.abc import Callable
+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
@@ -27,6 +34,8 @@ class ScopeMetrics:
         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
@@ -37,15 +46,38 @@ class ScopeMetrics:
             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._nested: set[ScopeMetrics] = set()
         self._timestamp: float = monotonic()
-        self._completed: Future[float] = get_event_loop().create_future()
+        self._finished: bool = False
+        self._loop: AbstractEventLoop = get_event_loop()
+        self._completed: Future[float] = self._loop.create_future()
+
+        if parent := parent:
+            parent._nested.add(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:
-        self._complete()  # ensure completion on deinit
+        assert self.is_completed, "Deinitializing not completed scope metrics"  # nosec: B101
 
     def __str__(self) -> str:
         return f"{self.label}[{self.identifier}]@[{self.trace_id}]"
@@ -113,8 +145,8 @@ class ScopeMetrics:
             self._metrics[metric_type] = metric
 
     @property
-    def completed(self) -> bool:
-        return self._completed.done() and all(nested.completed for nested in self._nested)
+    def is_completed(self) -> bool:
+        return self._completed.done() and all(nested.is_completed for nested in self._nested)
 
     @property
     def time(self) -> float:
@@ -131,24 +163,36 @@ class ScopeMetrics:
             return_exceptions=False,
         )
 
-    def _complete(self) -> None:
-        if self._completed.done():
-            return  # already completed
+    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"
 
+        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)
 
-    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
+        # notify parent about completion
+        if parent := self._parent:
+            parent._complete_if_able()
 
     def log(
         self,
@@ -178,29 +222,37 @@ class MetricsContext:
         *,
         trace_id: str | None = None,
         logger: Logger | None = None,
+        completion: Callable[[ScopeMetrics], Coroutine[None, None, None]]
+        | Callable[[ScopeMetrics], 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))
+        current: ScopeMetrics
+        try:  # check for current scope context
+            current = cls._context.get()
 
-            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
+        except LookupError:
+            # create metrics scope when missing yet
             return cls(
                 ScopeMetrics(
                     trace_id=trace_id,
                     scope=name,
                     logger=logger,
+                    parent=None,
+                    completion=completion,
                 )
             )
 
+        # or create nested metrics otherwise
+        return cls(
+            ScopeMetrics(
+                trace_id=trace_id,
+                scope=name,
+                logger=logger or current._logger,  # pyright: ignore[reportPrivateUsage]
+                parent=current,
+                completion=completion,
+            )
+        )
+
     @classmethod
     def record[Metric: State](
         cls,
@@ -320,15 +372,12 @@ class MetricsContext:
     ) -> 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
+            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._started = monotonic()
 
     def __exit__(
         self,
@@ -337,8 +386,8 @@ class MetricsContext:
         exc_tb: TracebackType | None,
     ) -> None:
         assert (  # nosec: B101
-            self._token is not None and self._started is not None and self._finished is None
+            self._token is not None
         ), "Unbalanced MetricsContext context enter/exit"
-        self._finished = monotonic()
         MetricsContext._context.reset(self._token)
+        self._metrics._finish()  # pyright: ignore[reportPrivateUsage]
         self._token = None

haiway/utils/queue.py

@@ -3,8 +3,6 @@ from collections import deque
 from collections.abc import AsyncIterator
 from typing import Self
 
-from haiway.utils.immutable import freeze
-
 __all__ = [
     "AsyncQueue",
 ]
@@ -18,20 +16,19 @@ class AsyncQueue[Element](AsyncIterator[Element]):
 
     def __init__(
         self,
+        *elements: Element,
         loop: AbstractEventLoop | None = None,
     ) -> None:
         self._loop: AbstractEventLoop = loop or get_running_loop()
-        self._queue: deque[Element] = deque()
+        self._queue: deque[Element] = deque(elements)
         self._waiting: Future[Element] | None = None
         self._finish_reason: BaseException | None = None
 
-        freeze(self)
-
     def __del__(self) -> None:
         self.finish()
 
     @property
-    def finished(self) -> bool:
+    def is_finished(self) -> bool:
         return self._finish_reason is not None
 
     def enqueue(
@@ -40,7 +37,7 @@ class AsyncQueue[Element](AsyncIterator[Element]):
         /,
         *elements: Element,
     ) -> None:
-        if self.finished:
+        if self.is_finished:
             raise RuntimeError("AsyncQueue is already finished")
 
         if self._waiting is not None and not self._waiting.done():
@@ -55,7 +52,7 @@ class AsyncQueue[Element](AsyncIterator[Element]):
         self,
         exception: BaseException | None = None,
     ) -> None:
-        if self.finished:
+        if self.is_finished:
             return  # already finished, ignore
 
         self._finish_reason = exception or StopAsyncIteration()
@@ -70,7 +67,7 @@ class AsyncQueue[Element](AsyncIterator[Element]):
         return self
 
     async def __anext__(self) -> Element:
-        assert self._waiting is None, "Only a single queue iterator is supported!"  # nosec: B101
+        assert self._waiting is None, "Only a single queue consumer is supported!"  # nosec: B101
 
         if self._queue:  # check the queue, let it finish
             return self._queue.popleft()

{haiway-0.3.2.dist-info → haiway-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: haiway
-Version: 0.3.2
+Version: 0.4.0
 Summary: Framework for dependency injection and state management within structured concurrency model.
 Maintainer-email: Kacper Kaliński <kacper.kalinski@miquido.com>
 License: MIT License

{haiway-0.3.2.dist-info → haiway-0.4.0.dist-info}/RECORD

@@ -1,9 +1,9 @@
 haiway/__init__.py,sha256=hLc3-FDmNQEV4r-RLOiGjWtYSk7krU8vRMBaZyRU08g,1267
 haiway/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haiway/context/__init__.py,sha256=21Y3zvRo1bHASZD6B_FNkU28k1-g88RdmUyqxvYXJxg,336
-haiway/context/access.py,sha256=E9aIC1RUupKk2LXik5qb13oHpW2s_tNQC0UxGcO9xr0,12761
+haiway/context/access.py,sha256=zPQcQBp5XMlNuszbxhtzi-5mNpDLpZ6PT-gVFDBH0dA,14115
 haiway/context/disposables.py,sha256=VQX9jVo1pjqkmOYzWpsbYyF45y0XtpjorIIaeMBCwTU,1771
-haiway/context/metrics.py,sha256=TGZzNrmGPolAwVb1mBEF10mXRyVyBq-MHTucAokNuhY,9105
+haiway/context/metrics.py,sha256=R6BIH3Pwtm_fQqPZmTei3nUGNDAZEnvUKaVxwjWxgS0,10755
 haiway/context/state.py,sha256=GxGwPQTK8FdSprBd83lQbA9veubp0o93_1Yk3gb7HMc,3000
 haiway/context/tasks.py,sha256=xXtXIUwXOra0EePTdkcEbMOmpWwFcO3hCRfR_IfvAHk,1978
 haiway/context/types.py,sha256=VvJA7wAPZ3ISpgyThVguioYUXqhHf0XkPfRd0M1ERiQ,142
@@ -28,9 +28,9 @@ haiway/utils/immutable.py,sha256=K34ZIMzbkpgkHKH-KF73plEbXExsajNRkRTYp9nJEf4,620
 haiway/utils/logs.py,sha256=oDsc1ZdqKDjlTlctLbDcp9iX98Acr-1tdw-Pyg3DElo,1577
 haiway/utils/mimic.py,sha256=BkVjTVP2TxxC8GChPGyDV6UXVwJmiRiSWeOYZNZFHxs,1828
 haiway/utils/noop.py,sha256=qgbZlOKWY6_23Zs43OLukK2HagIQKRyR04zrFVm5rWI,344
-haiway/utils/queue.py,sha256=WGW8kSusIwRYHsYRIKD2CaqhhC1pUtVgtNHFDXDtYrw,2443
-haiway-0.3.2.dist-info/LICENSE,sha256=GehQEW_I1pkmxkkj3NEa7rCTQKYBn7vTPabpDYJlRuo,1063
-haiway-0.3.2.dist-info/METADATA,sha256=Ze8RmuhIj2Ck3SVea8h3VbaTSf43_Y57oISLZKytNOA,3872
-haiway-0.3.2.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-haiway-0.3.2.dist-info/top_level.txt,sha256=_LdXVLzUzgkvAGQnQJj5kQfoFhpPW6EF4Kj9NapniLg,7
-haiway-0.3.2.dist-info/RECORD,,
+haiway/utils/queue.py,sha256=7gLpL07E4K_FnP1AygmpNnBpwfS5kgnw_6wakuRkmw4,2423
+haiway-0.4.0.dist-info/LICENSE,sha256=GehQEW_I1pkmxkkj3NEa7rCTQKYBn7vTPabpDYJlRuo,1063
+haiway-0.4.0.dist-info/METADATA,sha256=7inuv-8w44WoEhdlvmdPJkrbb-uz_pQ85KCpypIGf7w,3872
+haiway-0.4.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+haiway-0.4.0.dist-info/top_level.txt,sha256=_LdXVLzUzgkvAGQnQJj5kQfoFhpPW6EF4Kj9NapniLg,7
+haiway-0.4.0.dist-info/RECORD,,