haiway 0.3.2__tar.gz → 0.4.0__tar.gz

{haiway-0.3.2/src/haiway.egg-info → haiway-0.4.0}/PKG-INFO

@@ -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 → haiway-0.4.0}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "haiway"
 description = "Framework for dependency injection and state management within structured concurrency model."
-version = "0.3.2"
+version = "0.4.0"
 readme = "README.md"
 maintainers = [
   { name = "Kacper Kaliński", email = "kacper.kalinski@miquido.com" },

{haiway-0.3.2 → haiway-0.4.0}/src/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-0.3.2 → haiway-0.4.0}/src/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-0.3.2 → haiway-0.4.0}/src/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 → haiway-0.4.0/src/haiway.egg-info}/PKG-INFO

@@ -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 → haiway-0.4.0}/src/haiway.egg-info/SOURCES.txt

@@ -42,4 +42,5 @@ tests/test_auto_retry.py
 tests/test_cache.py
 tests/test_context.py
 tests/test_state.py
+tests/test_streaming.py
 tests/test_timeout.py

{haiway-0.3.2 → haiway-0.4.0}/tests/test_context.py

@@ -1,3 +1,5 @@
+from asyncio import get_running_loop
+
 from haiway import MissingContext, ScopeMetrics, State, ctx
 from pytest import mark, raises
 
@@ -64,38 +66,92 @@ async def test_exceptions_are_propagated():
 
 @mark.asyncio
 async def test_completions_are_called_according_to_context_exits():
+    completion_future = get_running_loop().create_future()
+    nested_completion_future = get_running_loop().create_future()
+    executions: int = 0
+
+    def completion(metrics: ScopeMetrics):
+        nonlocal executions
+        executions += 1
+        completion_future.set_result(())
+
+    def nested_completion(metrics: ScopeMetrics):
+        nonlocal executions
+        executions += 1
+        nested_completion_future.set_result(())
+
+    async with ctx.scope("outer", completion=completion):
+        assert executions == 0
+
+        async with ctx.scope("inner", completion=nested_completion):
+            assert executions == 0
+
+        await nested_completion_future
+        assert executions == 1
+
+    await completion_future
+    assert executions == 2
+
+
+@mark.asyncio
+async def test_async_completions_are_called_according_to_context_exits():
+    completion_future = get_running_loop().create_future()
+    nested_completion_future = get_running_loop().create_future()
     executions: int = 0
 
     async def completion(metrics: ScopeMetrics):
         nonlocal executions
         executions += 1
+        completion_future.set_result(())
+
+    async def nested_completion(metrics: ScopeMetrics):
+        nonlocal executions
+        executions += 1
+        nested_completion_future.set_result(())
 
     async with ctx.scope("outer", completion=completion):
         assert executions == 0
 
-        async with ctx.scope("inner", completion=completion):
+        async with ctx.scope("inner", completion=nested_completion):
             assert executions == 0
 
+        await nested_completion_future
         assert executions == 1
 
+    await completion_future
     assert executions == 2
 
 
 @mark.asyncio
 async def test_metrics_are_recorded_within_context():
-    def verify_example_metrics(state: str):
-        async def completion(metrics: ScopeMetrics):
-            assert metrics.read(ExampleState, default=ExampleState()).state == state
+    completion_future = get_running_loop().create_future()
+    nested_completion_future = get_running_loop().create_future()
+    metric: ExampleState = ExampleState()
+    nested_metric: ExampleState = ExampleState()
+
+    async def completion(metrics: ScopeMetrics):
+        nonlocal metric
+        metric = metrics.read(ExampleState, default=ExampleState())
+        completion_future.set_result(())
 
-        return completion
+    async def nested_completion(metrics: ScopeMetrics):
+        nonlocal nested_metric
+        nested_metric = metrics.read(ExampleState, default=ExampleState())
+        nested_completion_future.set_result(())
 
-    async with ctx.scope("outer", completion=verify_example_metrics("outer-in-out")):
+    async with ctx.scope("outer", completion=completion):
         ctx.record(ExampleState(state="outer-in"))
 
-        async with ctx.scope("inner", completion=verify_example_metrics("inner")):
+        async with ctx.scope("inner", completion=nested_completion):
             ctx.record(ExampleState(state="inner"))
 
+        await nested_completion_future
+        assert nested_metric.state == "inner"
+
         ctx.record(
             ExampleState(state="-out"),
             merge=lambda lhs, rhs: ExampleState(state=lhs.state + rhs.state),
         )
+
+    await completion_future
+    assert metric.state == "outer-in-out"

haiway-0.4.0/tests/test_streaming.py

@@ -0,0 +1,160 @@
+from asyncio import CancelledError, get_running_loop, sleep
+from collections.abc import AsyncGenerator, AsyncIterator
+
+from haiway import ctx
+from haiway.context.metrics import ScopeMetrics
+from haiway.state.structure import State
+from pytest import mark, raises
+
+
+class FakeException(Exception):
+    pass
+
+
+@mark.asyncio
+async def test_fails_when_generator_fails():
+    async def generator(value: int) -> AsyncGenerator[int, None]:
+        yield value
+        raise FakeException()
+
+    elements: int = 0
+    with raises(FakeException):
+        async for _ in ctx.stream(generator, 42):
+            elements += 1
+
+    assert elements == 1
+
+
+@mark.asyncio
+async def test_cancels_when_iteration_cancels():
+    async def generator(value: int) -> AsyncGenerator[int, None]:
+        await sleep(0)
+        yield value
+
+    elements: int = 0
+    with raises(CancelledError):
+        ctx.cancel()
+        async for _ in ctx.stream(generator, 42):
+            elements += 1
+
+    assert elements == 0
+
+
+@mark.asyncio
+async def test_ends_when_generator_ends():
+    async def generator(value: int) -> AsyncGenerator[int, None]:
+        yield value
+
+    elements: int = 0
+    async for _ in ctx.stream(generator, 42):
+        elements += 1
+
+    assert elements == 1
+
+
+@mark.asyncio
+async def test_delivers_updates_when_generating():
+    async def generator(value: int) -> AsyncGenerator[int, None]:
+        for i in range(0, value):
+            yield i
+
+    elements: list[int] = []
+
+    async for element in ctx.stream(generator, 10):
+        elements.append(element)
+
+    assert elements == list(range(0, 10))
+
+
+@mark.asyncio
+async def test_streaming_context_variables_access_is_preserved():
+    class TestState(State):
+        value: int = 42
+        other: str = "other"
+
+    async def generator(value: int) -> AsyncGenerator[TestState, None]:
+        yield ctx.state(TestState)
+        with ctx.scope("nested", ctx.state(TestState).updated(value=value)):
+            yield ctx.state(TestState)
+
+    stream: AsyncIterator[TestState]
+    async with ctx.scope("test", TestState(value=42)):
+        elements: list[TestState] = []
+
+        stream = ctx.stream(generator, 10)
+
+    async for element in stream:
+        elements.append(element)
+
+    assert elements == [
+        TestState(value=42),
+        TestState(value=10),
+    ]
+
+
+@mark.asyncio
+async def test_nested_streaming_streams_correctly():
+    class TestState(State):
+        value: int = 42
+        other: str = "other"
+
+    async def inner(value: int) -> AsyncGenerator[TestState, None]:
+        yield ctx.state(TestState)
+        with ctx.scope("inner", ctx.state(TestState).updated(value=value, other="inner")):
+            yield ctx.state(TestState)
+
+    async def outer(value: int) -> AsyncGenerator[TestState, None]:
+        yield ctx.state(TestState)
+        with ctx.scope("outer", ctx.state(TestState).updated(other="outer")):
+            async for item in ctx.stream(inner, value):
+                yield item
+
+    stream: AsyncIterator[TestState]
+    async with ctx.scope("test", TestState(value=42)):
+        elements: list[TestState] = []
+
+        stream = ctx.stream(outer, 10)
+
+    async for element in stream:
+        elements.append(element)
+
+    assert elements == [
+        TestState(value=42),
+        TestState(value=42, other="outer"),
+        TestState(value=10, other="inner"),
+    ]
+
+
+@mark.asyncio
+async def test_streaming_context_completion_is_called_at_the_end_of_stream():
+    completion_future = get_running_loop().create_future()
+
+    class IterationMetric(State):
+        value: int = 0
+
+    metric: IterationMetric = IterationMetric()
+
+    def completion(metrics: ScopeMetrics):
+        nonlocal metric
+        metric = metrics.metrics(merge=lambda current, nested: nested)[0]  # pyright: ignore[reportAssignmentType]
+        completion_future.set_result(())
+
+    async def generator(value: int) -> AsyncGenerator[int, None]:
+        for i in range(0, value):
+            ctx.record(
+                IterationMetric(value=i),
+                merge=lambda lhs, rhs: IterationMetric(value=lhs.value + rhs.value),
+            )
+            yield i
+
+    stream: AsyncIterator[int]
+    async with ctx.scope("test", completion=completion):
+        elements: list[int] = []
+
+        stream = ctx.stream(generator, 10)
+
+    async for element in stream:
+        elements.append(element)
+
+    await completion_future
+    assert metric.value == 45