pyfect-dev 0.1.0__py3-none-any.whl

pyfect/stream/_sink_methods.py

@@ -0,0 +1,151 @@
+"""``Stream`` methods that feed the stream to a ``Sink``: ``run``,
+``transduce``, ``aggregate``, ``aggregate_within``, ``grouped_within``,
+``pipe_through``, ``tap_sink`` and ``peel``. The work is in
+``pyfect.stream._sinking``."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Never, overload
+
+from pyfect.stream._base import A_co, E_co, R_co, StreamBase
+
+if TYPE_CHECKING:
+    from pyfect._core import Effect
+    from pyfect.duration import Duration
+    from pyfect.option import Option
+    from pyfect.schedule import Schedule
+    from pyfect.scope import Scope
+    from pyfect.sink import Sink
+    from pyfect.stream._stream import Stream
+
+
+def _require_sink(sink: object) -> None:
+    from pyfect.sink._sink import require_sink
+
+    require_sink(sink, "sink")
+
+
+class SinkMethods(StreamBase[A_co, E_co, R_co]):
+    __slots__ = ()
+
+    def run[B, L, E2 = Never, R2 = Never](
+        self, sink: Sink[B, A_co, L, E2, R2]
+    ) -> Effect[B, E_co | E2, R_co | R2]:
+        """Run the stream into ``sink`` and succeed with the sink's value.
+        The stream's failures pass by the sink's error handlers."""
+        from pyfect.stream import _sinking
+
+        _require_sink(sink)
+        return _sinking.run(self, sink)
+
+    # A stream of ``Never`` has no items to leave over, so it takes any
+    # sink; the overload keeps pyright from reading ``A`` as ``Never``.
+    @overload
+    def transduce[B, E, R, E2 = Never, R2 = Never](
+        self: SinkMethods[Never, E, R], sink: Sink[B, Never, object, E2, R2]
+    ) -> Stream[B, E | E2, R | R2]: ...
+    @overload
+    def transduce[A, E, R, B, E2 = Never, R2 = Never](
+        self: SinkMethods[A, E, R], sink: Sink[B, A, A, E2, R2]
+    ) -> Stream[B, E | E2, R | R2]: ...
+    def transduce(self, sink: Sink[Any, Any, Any, Any, Any]) -> Stream[Any, Any, Any]:
+        """Run ``sink`` again and again on the stream, writing each value it
+        ends with; its leftovers start the next run. A stream failure ends
+        the current run first."""
+        from pyfect.stream import _sinking
+
+        _require_sink(sink)
+        return _sinking.transduce(self, sink)
+
+    @overload
+    def aggregate[B, E, R, E2 = Never, R2 = Never](
+        self: SinkMethods[Never, E, R], sink: Sink[B, Never, object, E2, R2]
+    ) -> Stream[B, E | E2, R | R2]: ...
+    @overload
+    def aggregate[A, E, R, B, E2 = Never, R2 = Never](
+        self: SinkMethods[A, E, R], sink: Sink[B, A, A, E2, R2]
+    ) -> Stream[B, E | E2, R | R2]: ...
+    def aggregate(self, sink: Sink[Any, Any, Any, Any, Any]) -> Stream[Any, Any, Any]:
+        """``transduce`` while the stream keeps running: a fiber pulls the
+        stream as ``sink`` asks for input, and each run of the sink ends
+        when the sink is done."""
+        from pyfect import schedule
+        from pyfect.stream import _sinking
+
+        _require_sink(sink)
+        return _sinking.aggregate_within(self, sink, schedule.forever)
+
+    @overload
+    def aggregate_within[B, E, R, EE = Never, RR = Never](
+        self: SinkMethods[Never, E, R],
+        sink: Sink[B, Never, object, EE, RR],
+        schedule: Schedule[object, Option[B], EE, RR],
+    ) -> Stream[B, E | EE, R | RR]: ...
+    @overload
+    def aggregate_within[A, E, R, B, EE = Never, RR = Never](
+        self: SinkMethods[A, E, R],
+        sink: Sink[B, A, A, EE, RR],
+        schedule: Schedule[object, Option[B], EE, RR],
+    ) -> Stream[B, E | EE, R | RR]: ...
+    def aggregate_within(
+        self,
+        sink: Sink[Any, Any, Any, Any, Any],
+        schedule: Schedule[Any, Any, Any, Any],
+    ) -> Stream[Any, Any, Any]:
+        """``aggregate``, also ending each run of ``sink`` when ``schedule``
+        says so. The schedule steps once input arrives (with the last value
+        written) and its clock leaves out the time between runs; a schedule
+        that stops ends the stream once the sink has drained its
+        leftovers."""
+        from pyfect.stream import _sinking
+
+        _require_sink(sink)
+        return _sinking.aggregate_within(self, sink, schedule)
+
+    def grouped_within(self, chunk_size: int, duration: Duration) -> Stream[list[A_co], E_co, R_co]:
+        """Group the items in lists of up to ``chunk_size`` (at least 1),
+        ending a group early once ``duration`` has passed since its first
+        item."""
+        from pyfect import schedule
+        from pyfect.sink import take
+        from pyfect.stream import _sinking
+
+        return _sinking.aggregate_within(self, take(max(chunk_size, 1)), schedule.spaced(duration))
+
+    def pipe_through[B, L, E2 = Never, R2 = Never](
+        self, sink: Sink[B, A_co, L, E2, R2]
+    ) -> Stream[L, E_co | E2, R_co | R2]:
+        """Run the stream into ``sink`` and write the sink's leftovers. The
+        stream's failures pass by the sink's error handlers."""
+        from pyfect.stream import _sinking
+
+        _require_sink(sink)
+        return _sinking.pipe_through(self, sink)
+
+    def tap_sink[E2 = Never, R2 = Never](
+        self, sink: Sink[object, A_co, object, E2, R2]
+    ) -> Stream[A_co, E_co | E2, R_co | R2]:
+        """Send each chunk to ``sink`` too, as it is written. The stream
+        waits for the sink to take each chunk; once the sink is done it
+        goes on alone, and a sink failure fails the stream."""
+        from pyfect.stream import _sinking
+
+        _require_sink(sink)
+        return _sinking.tap_sink(self, sink)
+
+    @overload
+    def peel[B, E, R, E2 = Never, R2 = Never](
+        self: SinkMethods[Never, E, R], sink: Sink[B, Never, object, E2, R2]
+    ) -> Effect[tuple[B, Stream[Never, E]], E | E2, Scope | R | R2]: ...
+    @overload
+    def peel[A, E, R, B, E2 = Never, R2 = Never](
+        self: SinkMethods[A, E, R], sink: Sink[B, A, A, E2, R2]
+    ) -> Effect[tuple[B, Stream[A, E]], E | E2, Scope | R | R2]: ...
+    def peel(self, sink: Sink[Any, Any, Any, Any, Any]) -> Effect[Any, Any, Any]:
+        """Run ``sink`` on the start of the stream and succeed with its value
+        and the rest of the stream (the sink's leftovers first). The rest
+        reads the stream's pull, which lives in the current ``Scope``."""
+        from pyfect.stream import _sinking
+
+        _require_sink(sink)
+        return _sinking.peel(self, sink)

pyfect/stream/_sinking.py

@@ -0,0 +1,277 @@
+"""The work behind ``Stream``'s sink methods: ``run``, ``transduce``,
+``pipe_through``, ``tap_sink``, ``peel`` and ``aggregate_within``.
+
+A stream's own failures reach a sink as hidden defects and are restored
+once the sink has exited (``run_sink``), so no sink error handler sees
+them.
+"""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from typing import Any
+
+from pyfect import channel
+from pyfect._core import Effect, _FailCause, _Gen, _Succeed, _Suspend
+from pyfect.cause import Cause, fail
+from pyfect.channel._upstream import restore_upstream_failure, upstream_failure_to_defect
+from pyfect.clock import Clock, ClockService, sleep
+from pyfect.duration import from_nanos
+from pyfect.effect import never
+from pyfect.exit import Exit, Failure
+from pyfect.latch import Latch
+from pyfect.option import Option, Some
+from pyfect.pull import catch_done
+from pyfect.queue import Queue
+from pyfect.queue import _state as queue_state
+from pyfect.schedule import Schedule
+from pyfect.schedule._model import Continue, to_step
+from pyfect.scope import CloseableScope, fork_in, scoped_with
+from pyfect.sink._sink import AnySink, End
+from pyfect.sink._steps import ENDED
+from pyfect.stream._base import AnyStream
+from pyfect.stream._stream import Stream, transform_pull, transform_pull_bracket
+
+type AnyEffect = Effect[Any, Any, Any]
+type Chunk = tuple[Any, ...]
+type AnyStream_ = Stream[Any, Any, Any]
+type AnyEnd = End[Any, Any]
+type Channel_ = channel.Channel[Any, Any, Any, Any, Any, Any, Any]
+
+
+def run_sink(sink: AnySink, upstream: AnyEffect, scope: CloseableScope) -> AnyEffect:
+    """Run ``sink`` on ``upstream``; upstream failures bypass the sink's
+    error handlers and are restored once the sink has exited."""
+    return restore_upstream_failure(sink._transform(upstream_failure_to_defect(upstream), scope))
+
+
+def run(self: AnyStream, sink: AnySink) -> AnyEffect:
+    def consume(scope: CloseableScope) -> AnyEffect:
+        built = self._channel.to_pull_scoped(scope)
+        return built.flat_map(lambda pull: run_sink(sink, pull, scope)).map(_value)
+
+    return scoped_with(consume)
+
+
+def _value(end: AnyEnd) -> Any:
+    return end[0]
+
+
+def transduce(self: AnyStream, sink: AnySink) -> AnyStream_:
+    def build(upstream: AnyEffect, scope: CloseableScope) -> AnyEffect:
+        stopped: AnyEffect | None = None
+        leftover: Chunk | None = None
+
+        def next_chunk() -> AnyEffect:
+            nonlocal leftover
+            if leftover is not None:
+                chunk, leftover = leftover, None
+                return _Succeed(chunk)
+            return upstream
+
+        def stop(error: object) -> AnyEffect:
+            nonlocal stopped
+            stopped = _FailCause(fail(error))
+            return ENDED
+
+        fed = _Suspend(next_chunk).catch_all(stop)
+
+        def emit(end: AnyEnd) -> Chunk:
+            nonlocal leftover
+            leftover = end[1]
+            return (end[0],)
+
+        pull = _Suspend(lambda: sink._transform(fed, scope)).map(emit)
+        return _Succeed(_Suspend(lambda: pull if stopped is None else stopped))
+
+    return transform_pull(self, build)
+
+
+def pipe_through(self: AnyStream, sink: AnySink) -> AnyStream_:
+    from pyfect.sink import to_channel
+
+    def leftovers(end: AnyEnd) -> Any:
+        return channel.empty if end[1] is None else channel.succeed(end[1])
+
+    piped: Any = self._channel.pipe_to_or_fail(to_channel(sink)).concat_with(leftovers)
+    return Stream(piped)
+
+
+def tap_sink(self: AnyStream, sink: AnySink) -> AnyStream_:
+    def build(pull: AnyEffect, _scope: CloseableScope, forked: CloseableScope) -> AnyEffect:
+        upstream_latch = Latch()
+        sink_latch = Latch()
+        chunk: Chunk | None = None
+        sink_cause: Cause[Any] | None = None
+        sink_done = False
+        stream_done = False
+
+        def hand_over() -> AnyEffect:
+            nonlocal chunk
+            if chunk is None:
+                return ENDED
+            taken, chunk = chunk, None
+            if not stream_done:
+                upstream_latch.close_unsafe()
+            return sink_latch.open().as_(taken)
+
+        sink_upstream = upstream_latch.when_open(_Suspend(hand_over))
+
+        def sink_exited(ex: Exit[Any, Any]) -> AnyEffect:
+            nonlocal sink_done, sink_cause
+            sink_done = True
+            if isinstance(ex, Failure):
+                sink_cause = ex.cause
+            return sink_latch.open()
+
+        running = _Suspend(lambda: sink._transform(sink_upstream, forked)).on_exit(sink_exited)
+
+        def offer(next_chunk: Chunk) -> AnyEffect:
+            nonlocal chunk
+            if sink_cause is not None:
+                return _FailCause(sink_cause)
+            if sink_done:
+                return _Succeed(next_chunk)
+            chunk = next_chunk
+            sink_latch.close_unsafe()
+            upstream_latch.open_unsafe()
+            return sink_latch.wait().as_(next_chunk)
+
+        def stream_ended(_: object) -> AnyEffect:
+            nonlocal stream_done
+            if sink_cause is not None:
+                return _FailCause(sink_cause)
+            if sink_done:
+                return ENDED
+            stream_done = True
+            sink_latch.close_unsafe()
+            upstream_latch.open_unsafe()
+            return sink_latch.wait().and_then(
+                _Suspend(lambda: ENDED if sink_cause is None else _FailCause(sink_cause))
+            )
+
+        pull_and_offer = catch_done(pull.flat_map(offer), stream_ended)
+
+        def next_pull() -> AnyEffect:
+            if sink_cause is not None:
+                return _FailCause(sink_cause)
+            return pull if sink_done else pull_and_offer
+
+        return fork_in(running, forked).as_(_Suspend(next_pull))
+
+    return transform_pull_bracket(self, build)
+
+
+def peel(self: AnyStream, sink: AnySink) -> AnyEffect:
+    from pyfect.stream._constructors import empty, from_iterable
+    from pyfect.stream._stream import from_pull
+
+    def body() -> Generator[AnyEffect, Any, tuple[Any, AnyStream_]]:
+        failed = False
+        original = yield self._channel.to_pull()
+
+        def mark(cause: Cause[Any]) -> AnyEffect:
+            nonlocal failed
+            failed = True
+            return _FailCause(cause)
+
+        pull = original.catch_cause(mark)
+        value, leftover = yield scoped_with(lambda scope: run_sink(sink, pull, scope))
+        rest: AnyStream_ = empty if failed else from_pull(_Succeed(original))
+        if leftover is not None:
+            rest = from_iterable(leftover).concat(rest)
+        return (value, rest)
+
+    return _Gen(body)
+
+
+_SCHEDULE_STEP = object()
+"""What the schedule puts in the buffer to end the current aggregation."""
+
+
+def aggregate_within(
+    self: AnyStream, sink: AnySink, policy: Schedule[Any, Any, Any, Any]
+) -> AnyStream_:
+    def transform(_upstream: Any, scope: CloseableScope, forked: CloseableScope) -> AnyEffect:
+        def body() -> Generator[AnyEffect, Any, AnyEffect]:
+            pull = yield self._channel.to_pull_scoped(scope)
+            pull_latch = Latch()
+            buffer: Queue[Any, Any] = Queue(0)
+
+            def offer(chunk: Chunk) -> AnyEffect:
+                pull_latch.close_unsafe()
+                return buffer.offer(chunk)
+
+            producer = pull_latch.when_open(pull).flat_map(offer).forever()
+            yield fork_in(producer.catch_cause(buffer.fail_cause), forked)
+
+            # The schedule's clock leaves out the time between aggregations.
+            clock: ClockService = yield Clock()
+            step = yield to_step(policy)
+            has_input = Latch()
+            schedule_time = clock.now_nanos_unsafe()
+            opened_at = 0
+            last_output: Option[Any] = None
+            leftover: Chunk | None = None
+
+            def on_input(chunk: Chunk) -> AnyEffect:
+                nonlocal opened_at
+                if has_input.open_unsafe():
+                    opened_at = clock.now_nanos_unsafe()
+                return _Succeed(chunk)
+
+            def decided(decision: object) -> AnyEffect:
+                if isinstance(decision, Continue):
+                    return sleep(decision.delay).and_then(buffer.offer(_SCHEDULE_STEP))
+                # End input so the sink drains its leftovers without pulling
+                # more upstream.
+                return buffer.end()
+
+            step_to_buffer = (
+                has_input.wait()
+                .flat_map(lambda _: step(from_nanos(schedule_time), last_output))
+                .flat_map(decided)
+                .and_then(never)
+            )
+
+            def take_input(item: Any) -> AnyEffect:
+                return ENDED if item is _SCHEDULE_STEP else on_input(item)
+
+            def next_input() -> AnyEffect:
+                nonlocal leftover
+                if leftover is not None:
+                    chunk, leftover = leftover, None
+                    return on_input(chunk)
+                if isinstance(buffer._state, queue_state.Open):
+                    pull_latch.open_unsafe()
+                return buffer.take().flat_map(take_input)
+
+            sink_upstream = _Suspend(next_input)
+
+            def emit(end: AnyEnd) -> AnyEffect:
+                nonlocal schedule_time, last_output, leftover
+                if has_input.is_open():
+                    schedule_time += clock.now_nanos_unsafe() - opened_at
+                elif isinstance(buffer._state, queue_state.Done):
+                    return ENDED
+                last_output = Some(end[0])
+                leftover = end[1]
+                return _Succeed((end[0],))
+
+            def next_pull() -> AnyEffect:
+                state = buffer._state
+                if isinstance(state, queue_state.Done) and leftover is None:
+                    return _FailCause(state.cause)
+                has_input.close_unsafe()
+                aggregated = _Suspend(lambda: run_sink(sink, sink_upstream, scope)).flat_map(emit)
+                if isinstance(state, queue_state.Open):
+                    return aggregated.race_first(step_to_buffer)
+                return aggregated
+
+            return _Suspend(next_pull)
+
+        return _Gen(body)
+
+    built: Channel_ = channel.from_transform_bracket(transform)
+    return Stream(built)

pyfect/stream/_sources.py

@@ -0,0 +1,124 @@
+"""Streams fed from outside: callbacks offering to a queue, queues,
+pubsubs, schedules, ticks and async iterables."""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterable, Callable
+from typing import Any, Never, overload
+
+from pyfect import channel
+from pyfect._core import Effect, _Succeed, _Suspend, _Sync
+from pyfect.channel._base import end_with
+from pyfect.duration import Duration, to_seconds
+from pyfect.errors import Done
+from pyfect.pubsub import PubSub, Subscription
+from pyfect.queue import Dequeue, Queue, Strategy
+from pyfect.schedule import Schedule
+from pyfect.schedule._model import Metadata, Stop, to_step_with_metadata
+from pyfect.scope import Scope
+from pyfect.stream._stream import Stream, from_pull
+from pyfect.take import Take
+
+type AnyPull = Effect[Any, Any, Any]
+
+
+@overload
+def callback[A, E = Never](
+    f: Callable[[Queue[A, E | Done]], Effect[object, E, Scope]],
+    *,
+    buffer_size: int | None = None,
+    strategy: Strategy = "suspend",
+) -> Stream[A, E]: ...
+@overload
+def callback[A, E = Never, R = Never](
+    f: Callable[[Queue[A, E | Done]], Effect[object, E, R | Scope]],
+    *,
+    buffer_size: int | None = None,
+    strategy: Strategy = "suspend",
+) -> Stream[A, E, R]: ...
+def callback(
+    f: Callable[[Queue[Any, Any]], Effect[object, Any, Any]],
+    *,
+    buffer_size: int | None = None,
+    strategy: Strategy = "suspend",
+) -> Stream[Any, Any, Any]:
+    """A stream of what ``f`` offers to a queue. ``f`` runs on a fiber, with
+    the stream's ``Scope``, when the stream starts; the stream ends when the
+    queue ends and fails with ``f``'s failure or the queue's. With a
+    ``buffer_size`` the queue is bounded and ``strategy`` says what a full
+    one does (``"suspend"`` waits: backpressure)."""
+    chunks: Any = channel.callback_array(f, buffer_size=buffer_size, strategy=strategy)
+    return Stream(chunks)
+
+
+@overload
+def from_queue[A](queue: Dequeue[A, Done]) -> Stream[A]: ...
+@overload
+def from_queue[A, E](queue: Dequeue[A, E | Done]) -> Stream[A, E]: ...
+def from_queue(queue: Dequeue[Any, Any]) -> Stream[Any, Any]:
+    """A stream taking everything waiting in ``queue`` at each pull, until
+    it ends (or fails)."""
+    return Stream(channel.from_queue_array(queue))
+
+
+def from_subscription[A](subscription: Subscription[A]) -> Stream[A]:
+    """A stream of the messages of ``subscription``; it ends when the
+    pubsub shuts down."""
+    return Stream(channel.from_subscription_array(subscription))
+
+
+def from_pubsub[A](pubsub: PubSub[A]) -> Stream[A]:
+    """A stream of what is published to ``pubsub`` after the stream starts
+    (and what it replays); it ends when the pubsub shuts down."""
+    return Stream(channel.from_pubsub_array(pubsub))
+
+
+def from_pubsub_take[A, E](pubsub: PubSub[Take[A, E, None]]) -> Stream[A, E]:
+    """A stream replaying the ``Take`` values published to ``pubsub``: the
+    chunks, then the failure or the end."""
+    return Stream(channel.from_pubsub_take(pubsub))
+
+
+def from_schedule[O, E = Never, R = Never](schedule: Schedule[O, object, E, R]) -> Stream[O, E, R]:
+    """A stream of the outputs of ``schedule``, each after its delay; it
+    ends when the schedule stops."""
+
+    def with_step(step: Callable[[None], Effect[Metadata[O, Any] | Stop[O], E, R]]) -> AnyPull:
+        def written(decision: Metadata[O, Any] | Stop[O]) -> AnyPull:
+            return end_with(None) if isinstance(decision, Stop) else _Succeed((decision.output,))
+
+        return step(None).flat_map(written)
+
+    return from_pull(to_step_with_metadata(schedule).map(with_step))
+
+
+def tick(interval: Duration) -> Stream[None]:
+    """An endless stream of ``None``: the first at once, then one after
+    each ``interval``."""
+    seconds = to_seconds(interval)
+    tock: Effect[tuple[None]] = _Succeed((None,))
+
+    def make_pull() -> AnyPull:
+        delayed = tock.delay(seconds)
+        first = True
+
+        def pull() -> AnyPull:
+            nonlocal first
+            if first:
+                first = False
+                return tock
+            return delayed
+
+        return _Suspend(pull)
+
+    return from_pull(_Sync(make_pull))
+
+
+def from_async_iterable[A, E](
+    iterable: AsyncIterable[A], on_error: Callable[[Exception], E]
+) -> Stream[A, E]:
+    """A stream of what ``iterable`` yields, failing with
+    ``on_error(exception)`` if it raises. The iterator is closed when the
+    run ends. Needs an asyncio loop (``run_async``)."""
+    return Stream(channel.from_async_iterable_array(iterable, on_error))