sink

from pyfect import sink, stream

fx.run_sync(stream.range(0, 5).run(sink.sum()))  # 10

die(defect) -> Sink[Never]

def die(defect: BaseException) -> Sink[Never]:
    """A sink that reads nothing and dies with ``defect``."""
    return from_effect_end(_effect.die(defect))

fail(error) -> Sink[Never, object, Never, E]

def fail[E](error: E) -> Sink[Never, object, Never, E]:
    """A sink that reads nothing and fails with ``error``."""
    return from_effect_end(_effect.fail(error))

fail_cause(cause) -> Sink[Never, object, Never, E]

def fail_cause[E](cause: Cause[E]) -> Sink[Never, object, Never, E]:
    """A sink that reads nothing and fails with ``cause``."""
    return from_effect_end(_effect.fail_cause(cause))

fail_cause_sync(thunk) -> Sink[Never, object, Never, E]

def fail_cause_sync[E](thunk: Callable[[], Cause[E]]) -> Sink[Never, object, Never, E]:
    """A sink that reads nothing and fails with the cause ``thunk()``."""
    return from_effect_end(_effect.fail_cause_sync(thunk))

fail_sync(thunk) -> Sink[Never, object, Never, E]

def fail_sync[E](thunk: Callable[[], E]) -> Sink[Never, object, Never, E]:
    """A sink that reads nothing and fails with ``thunk()``."""
    return from_effect_end(_effect.fail_sync(thunk))

from_channel(channel) -> Sink[A, In, L, E, R]

def from_channel[A, In, L = Never, E = Never, R = Never](
    channel: Channel[Never, E, End[A, L], tuple[In, ...], Never, None, R],
) -> Sink[A, In, L, E, R]:
    """The sink that feeds its input to ``channel`` and ends with the
    ``End`` the channel ends with."""
    require_channel(channel, "channel")

    def transform(upstream: Upstream[Any], scope: CloseableScope) -> AnyEffect:
        drained = channel._transform(upstream, scope).flat_map(lambda pull: PullLoop(pull, _ignore))
        return catch_done(drained, _Succeed)

    return Sink(transform)

from_effect(eff) -> Sink[A, object, Never, E, R]

def from_effect[A, E = Never, R = Never](eff: Effect[A, E, R]) -> Sink[A, object, Never, E, R]:
    """A sink that reads nothing and succeeds with ``eff``'s value."""
    require_effect(eff, "eff")
    return from_effect_end(eff.map(lambda value: (value, None)))

from_effect_end(eff) -> Sink[A, object, L, E, R]

def from_effect_end[A, L = Never, E = Never, R = Never](
    eff: Effect[End[A, L], E, R],
) -> Sink[A, object, L, E, R]:
    """A sink that reads nothing and ends with what ``eff`` succeeds
    with."""
    require_effect(eff, "eff")
    return Sink(lambda _upstream, _scope: eff)

from_pubsub(pubsub) -> Sink[None, A]

def from_pubsub[A](pubsub: PubSub[A]) -> Sink[None, A]:
    """A sink that publishes its input to ``pubsub``."""
    return for_each_array(pubsub.publish_all)

from_queue(queue) -> Sink[None, A]

def from_queue[A](queue: Enqueue[A, Done]) -> Sink[None, A]:
    """A sink that offers its input to ``queue`` and ends the queue when the
    input ends."""

    def transform(upstream: Upstream[A], _scope: CloseableScope) -> AnyEffect:
        offered = PullLoop(upstream.flat_map(queue.offer_all), _ignore)
        return catch_done(offered, lambda _: queue.end().and_then(_END_VOID))

    return Sink(transform)

from_transform(transform) -> Sink[A, In, L, E, R]

def from_transform[In, A, L = Never, E = Never, R = Never](
    transform: Callable[[Upstream[In], CloseableScope], Effect[End[A, L], E, R]],
) -> Sink[A, In, L, E, R]:
    """A sink from ``transform(upstream, scope)``: an effect that pulls
    non-empty tuples from ``upstream`` (which ends with ``Done()``) and
    succeeds with the sink's ``End``. Resources go in ``scope``."""
    return Sink(transform)

make(f) -> AnySink

sink.make(lambda s: s.run_sum())

def make(f: Callable[[Any], AnyEffect]) -> AnySink:
    """A sink from a stream pipeline: ``f`` gets the sink's input as a
    ``Stream`` and returns the effect whose value is the sink's (Effect's
    ``Sink.make<In>()(...)``). The effect runs with the sink's ``Scope``::

        sink.make(lambda s: s.run_sum())
    """
    from pyfect.stream import from_pull

    def transform(upstream: Upstream[Any], scope: CloseableScope) -> AnyEffect:
        run = f(from_pull(_Succeed(upstream))).map(lambda value: (value, None))
        return _Locally(run, lambda ctx: ctx.add(Scope, scope))

    return Sink(transform)

succeed(value, leftover=None) -> Sink[A, object, L]

def succeed[A, L = Never](value: A, leftover: tuple[L, ...] | None = None) -> Sink[A, object, L]:
    """A sink that reads nothing and succeeds with ``value``, leaving
    ``leftover`` over (an empty tuple is no leftover)."""
    end: End[A, L] = (value, leftover or None)
    return from_effect_end(_Succeed(end))

suspend(thunk) -> Sink[A, In, L, E, R]

def suspend[A, In, L, E, R](thunk: Callable[[], Sink[A, In, L, E, R]]) -> Sink[A, In, L, E, R]:
    """The sink ``thunk()`` returns, built each time it runs."""
    return Sink(lambda upstream, scope: thunk()._transform(upstream, scope))

sync(thunk) -> Sink[A]

def sync[A](thunk: Callable[[], A]) -> Sink[A]:
    """A sink that reads nothing and succeeds with ``thunk()``, called
    each time it runs."""
    return from_effect_end(_Sync(lambda: (thunk(), None)))

to_channel(sink) -> Channel[Never, E, End[A, L], tuple[In, ...], Never, None, R]

def to_channel[A, In, L, E, R](
    sink: Sink[A, In, L, E, R],
) -> Channel[Never, E, End[A, L], tuple[In, ...], Never, None, R]:
    """The channel that reads what the sink reads and ends with its
    ``End``."""
    require_sink(sink, "sink")

    def transform(upstream: Any, scope: CloseableScope) -> AnyEffect:
        return _Succeed(sink._transform(upstream, scope).flat_map(end_with))

    built: Channel[Any, Any, Any, Any, Any, Any, Any] = Channel(transform)
    return built

provide_context(ctx) -> _ProvideContext[RO]

def provide_context[RO](ctx: Context[RO]) -> _ProvideContext[RO]:
    """Provide every service of ``ctx`` to a sink and remove them from
    ``R``::

    sink.pipe(provide_context(ctx))
    """
    return _ProvideContext(ctx)

provide_service(tag, service) -> _ProvideService[S, T]

def provide_service[S, T](tag: TagLike[S, T], service: S) -> _ProvideService[S, T]:
    """Provide one service to a sink and remove it from ``R``::

    sink.pipe(provide_service(Database, InMemoryDb()))
    """
    return _ProvideService(tag_class(tag), service)

unwrap(eff) -> Sink_

def unwrap(eff: Effect[AnySink, Any, Any]) -> Sink_:
    """The sink ``eff`` succeeds with, built each time it runs. ``eff``
    runs with the sink's ``Scope``, so its resources live as long as the
    sink's run."""
    require_effect(eff, "eff")

    def transform(upstream: Upstream[Any], scope: CloseableScope) -> Effect[Any, Any, Any]:
        built: Effect[AnySink, Any, Any] = _Locally(eff, lambda ctx: ctx.add(Scope, scope))
        return built.flat_map(lambda sink: sink._transform(upstream, scope))

    return Sink(transform)

collect() -> Sink[list[In], In]

def collect[In]() -> Sink[list[In], In]:
    """A sink collecting its input in a list."""
    return reduce_array(list, _extend)

every(predicate) -> Sink[bool, In, In]

def every[In](predicate: Callable[[In], bool]) -> Sink[bool, In, In]:
    """``True`` if ``predicate`` holds for every item; stops at the first
    item it fails for."""
    return _fold_bool(True, predicate)

find(refinement) -> Sink[Option[Any], Any, Any]

def find(refinement: Callable[[Any], bool]) -> Sink[Option[Any], Any, Any]:
    """``Some`` of the first item ``predicate`` holds for (the rest of its
    chunk is left over), or ``None``."""

    def check(found: Option[Any], item: Any) -> Option[Any]:
        return Some(item) if refinement(item) else found

    return reduce_while(_nothing, _not_found, check)

find_effect(predicate) -> Sink[Option[In], In, In, E, R]

def find_effect[In, E = Never, R = Never](
    predicate: Callable[[In], Effect[bool, E, R]],
) -> Sink[Option[In], In, In, E, R]:
    """``find`` with an effectful ``predicate``."""

    def check(found: Option[In], item: In) -> Effect[Option[In], E, R]:
        return predicate(item).map(lambda ok: Some(item) if ok else found)

    return reduce_while_effect(_nothing, _not_found, check)

fold(initial, cont, f) -> Sink[S, In, In, E, R]

def fold[S, In, E = Never, R = Never](
    initial: Callable[[], S],
    cont: Callable[[S], bool],
    f: Callable[[S, In], Effect[S, E, R]],
) -> Sink[S, In, In, E, R]:
    """Fold each item into ``initial()`` with ``f`` while ``cont(state)``
    holds; once it fails, end with the state and the rest of the chunk as
    leftovers. ``cont`` is first asked after the first item."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        state = initial()

        def body() -> Generator[AnyEffect, Any, End[S, In]]:
            nonlocal state
            while True:
                chunk = yield upstream
                for i, item in enumerate(chunk):
                    state = yield f(state, item)
                    if not cont(state):
                        return (state, _rest(chunk, i + 1))

        return catch_done(_Gen(body), lambda _: _Sync(lambda: (state, None)))

    return Sink(transform)

fold_array(initial, cont, f) -> Sink[S, In, Never, E, R]

def fold_array[S, In, E = Never, R = Never](
    initial: Callable[[], S],
    cont: Callable[[S], bool],
    f: Callable[[S, tuple[In, ...]], Effect[S, E, R]],
) -> Sink[S, In, Never, E, R]:
    """``fold`` a chunk at a time; nothing is left over."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        state = initial()

        def step(chunk: tuple[In, ...]) -> AnyEffect:
            def keep(next_state: S) -> AnyEffect:
                nonlocal state
                state = next_state
                return VOID if cont(state) else ENDED

            return f(state, chunk).flat_map(keep)

        return read_until_done(upstream.flat_map(step), lambda: (state, None))

    return Sink(transform)

fold_until(initial, max, f) -> Sink[S, In, In, E, R]

def fold_until[S, In, E = Never, R = Never](
    initial: Callable[[], S], max: int, f: Callable[[S, In], Effect[S, E, R]]
) -> Sink[S, In, In, E, R]:
    """Fold at most ``max`` items into ``initial()`` with ``f``. With
    ``max`` 0 or less the sink reads nothing."""
    if max <= 0:
        return sync(initial)

    def start() -> tuple[S, int]:
        return (initial(), 0)

    def below_max(pair: tuple[S, int]) -> bool:
        return pair[1] < max

    def step(pair: tuple[S, int], item: In) -> Effect[tuple[S, int], E, R]:
        return f(pair[0], item).map(lambda state: (state, pair[1] + 1))

    return fold(start, below_max, step).map(lambda pair: pair[0])

head() -> Sink[Option[In], In, In]

def head[In]() -> Sink[Option[In], In, In]:
    """``Some`` of the first item (the rest of its chunk is left over), or
    ``None`` if there is none."""
    return reduce_while(_nothing, _not_found, _first)

last() -> Sink[Option[In], In]

def last[In]() -> Sink[Option[In], In]:
    """``Some`` of the last item, or ``None`` if there is none."""
    return reduce_array(_nothing, _last)

reduce(initial, f) -> Sink[S, In]

def reduce[S, In](initial: Callable[[], S], f: Callable[[S, In], S]) -> Sink[S, In]:
    """Reduce every item into ``initial()`` with ``f``."""

    def reduce_chunk(state: S, chunk: tuple[In, ...]) -> S:
        for item in chunk:
            state = f(state, item)
        return state

    return reduce_array(initial, reduce_chunk)

reduce_array(initial, f) -> Sink[S, In]

def reduce_array[S, In](
    initial: Callable[[], S], f: Callable[[S, tuple[In, ...]], S]
) -> Sink[S, In]:
    """Reduce every chunk into ``initial()`` with ``f``."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        state = initial()

        def step(chunk: tuple[In, ...]) -> None:
            nonlocal state
            state = f(state, chunk)

        return read_until_done(upstream.map(step), lambda: (state, None))

    return Sink(transform)

reduce_effect(initial, f) -> Sink[S, In, Never, E, R]

def reduce_effect[S, In, E = Never, R = Never](
    initial: Callable[[], S], f: Callable[[S, In], Effect[S, E, R]]
) -> Sink[S, In, Never, E, R]:
    """Reduce every item into ``initial()`` with ``f``."""
    reduced: Sink[S, In, Any, E, R] = reduce_while_effect(initial, lambda _: True, f)
    return reduced

reduce_while(initial, predicate, f) -> Sink[S, In, In]

def reduce_while[S, In](
    initial: Callable[[], S], predicate: Callable[[S], bool], f: Callable[[S, In], S]
) -> Sink[S, In, In]:
    """Reduce each item into ``initial()`` with ``f`` while
    ``predicate(state)`` holds, asking it of ``initial()`` first; the rest
    of the chunk is left over."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        state = initial()
        leftover: Chunk | None = None
        if not predicate(state):
            return _Succeed((state, None))

        def step(chunk: Chunk) -> AnyEffect:
            nonlocal state, leftover
            for i, item in enumerate(chunk):
                state = f(state, item)
                if not predicate(state):
                    leftover = _rest(chunk, i + 1)
                    return ENDED
            return VOID

        return read_until_done(upstream.flat_map(step), lambda: (state, leftover))

    return Sink(transform)

reduce_while_array(initial, predicate, f) -> Sink[S, In]

def reduce_while_array[S, In](
    initial: Callable[[], S], predicate: Callable[[S], bool], f: Callable[[S, tuple[In, ...]], S]
) -> Sink[S, In]:
    """``reduce_while`` a chunk at a time; nothing is left over."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        state = initial()
        if not predicate(state):
            return _Succeed((state, None))

        def step(chunk: tuple[In, ...]) -> AnyEffect:
            nonlocal state
            state = f(state, chunk)
            return VOID if predicate(state) else ENDED

        return read_until_done(upstream.flat_map(step), lambda: (state, None))

    return Sink(transform)

reduce_while_array_effect(initial, predicate, f) -> Sink[S, In, Never, E, R]

def reduce_while_array_effect[S, In, E = Never, R = Never](
    initial: Callable[[], S],
    predicate: Callable[[S], bool],
    f: Callable[[S, tuple[In, ...]], Effect[S, E, R]],
) -> Sink[S, In, Never, E, R]:
    """``reduce_while_effect`` a chunk at a time; nothing is left over."""

    def transform(upstream: Upstream[In], scope: CloseableScope) -> AnyEffect:
        state = initial()
        if not predicate(state):
            return _Succeed((state, None))
        return fold_array(lambda: state, predicate, f)._transform(upstream, scope)

    return Sink(transform)

reduce_while_effect(initial, predicate, f) -> Sink[S, In, In, E, R]

def reduce_while_effect[S, In, E = Never, R = Never](
    initial: Callable[[], S],
    predicate: Callable[[S], bool],
    f: Callable[[S, In], Effect[S, E, R]],
) -> Sink[S, In, In, E, R]:
    """Reduce each item into ``initial()`` with ``f`` while
    ``predicate(state)`` holds, asking it of ``initial()`` first; the rest
    of the chunk is left over."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        state = initial()
        leftover: Chunk | None = None
        if not predicate(state):
            return _Succeed((state, None))

        def step(chunk: Chunk) -> AnyEffect:
            def body() -> Generator[AnyEffect, Any, AnyEffect]:
                nonlocal state, leftover
                for i, item in enumerate(chunk):
                    state = yield f(state, item)
                    if not predicate(state):
                        leftover = _rest(chunk, i + 1)
                        return ENDED
                return VOID

            return _Gen(body).flatten()

        return read_until_done(upstream.flat_map(step), lambda: (state, leftover))

    return Sink(transform)

some(predicate) -> Sink[bool, In, In]

def some[In](predicate: Callable[[In], bool]) -> Sink[bool, In, In]:
    """``True`` if ``predicate`` holds for some item; stops at the first
    item it holds for."""
    return _fold_bool(False, predicate)

sum_() -> Sink[N, N]

def sum_[N: (int, float)]() -> Sink[N, N]:
    """A sink adding up its input (``0`` for none). Effect's ``Sink.sum``,
    a function here so the result has the input's type."""
    summed: Sink[Any, Any] = reduce_array(int, _add_up)
    return summed

for_each(f) -> Sink[None, In, Never, E, R]

def for_each[In, E = Never, R = Never](
    f: Callable[[In], Effect[object, E, R]],
) -> Sink[None, In, Never, E, R]:
    """A sink that runs ``f`` on each item of its input, in order."""

    def each(chunk: Chunk) -> AnyEffect:
        if len(chunk) == 1:
            return f(chunk[0])

        def body() -> Generator[AnyEffect, Any]:
            for item in chunk:
                yield f(item)

        return _Gen(body)

    return for_each_array(each)

for_each_array(f) -> Sink[None, In, Never, E, R]

def for_each_array[In, E = Never, R = Never](
    f: Callable[[tuple[In, ...]], Effect[object, E, R]],
) -> Sink[None, In, Never, E, R]:
    """A sink that runs ``f`` on each chunk of its input, in order."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        return _until_done(upstream.flat_map(f))

    return Sink(transform)

for_each_while(f) -> Sink[None, In, Never, E, R]

def for_each_while[In, E = Never, R = Never](
    f: Callable[[In], Effect[bool, E, R]],
) -> Sink[None, In, Never, E, R]:
    """A sink that runs ``f`` on each item until it succeeds with
    ``False``; the sink stops reading there."""

    def each(chunk: Chunk) -> Effect[bool, Any, Any]:
        def body() -> Generator[AnyEffect, Any, bool]:
            for item in chunk:
                if not (yield f(item)):
                    return False
            return True

        return _Gen(body)

    return for_each_while_array(each)

for_each_while_array(f) -> Sink[None, In, Never, E, R]

def for_each_while_array[In, E = Never, R = Never](
    f: Callable[[tuple[In, ...]], Effect[bool, E, R]],
) -> Sink[None, In, Never, E, R]:
    """A sink that runs ``f`` on each chunk until it succeeds with
    ``False``; the sink stops reading there."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        return _until_done(upstream.flat_map(f).flat_map(_go_on))

    return Sink(transform)

take(n) -> Sink[list[In], In, In]

def take[In](n: int) -> Sink[list[In], In, In]:
    """A sink taking the first ``n`` items (none for ``n`` 0 or less); the
    rest of the chunk that reaches ``n`` is left over."""
    count = max(n, 0)

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        taken: list[In] = []
        if count == 0:
            return _Succeed((taken, None))
        leftover: Chunk | None = None

        def step(chunk: tuple[In, ...]) -> AnyEffect:
            nonlocal leftover
            room = count - len(taken)
            if len(chunk) < room:
                taken.extend(chunk)
                return VOID
            taken.extend(chunk[:room])
            leftover = _rest(chunk, room)
            return ENDED

        return read_until_done(upstream.flat_map(step), lambda: (taken, leftover))

    return Sink(transform)

take_until(predicate) -> Sink[list[In], In, In]

def take_until[In](predicate: Callable[[In], bool]) -> Sink[list[In], In, In]:
    """A sink taking items up to and including the first one
    ``predicate`` holds for."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        taken: list[In] = []
        leftover: Chunk | None = None

        def step(chunk: tuple[In, ...]) -> AnyEffect:
            nonlocal leftover
            for i, item in enumerate(chunk):
                taken.append(item)
                if predicate(item):
                    leftover = _rest(chunk, i + 1)
                    return ENDED
            return VOID

        return read_until_done(upstream.flat_map(step), lambda: (taken, leftover))

    return Sink(transform)

take_until_effect(predicate) -> Sink[list[In], In, In, E, R]

def take_until_effect[In, E = Never, R = Never](
    predicate: Callable[[In], Effect[bool, E, R]],
) -> Sink[list[In], In, In, E, R]:
    """``take_until`` with an effectful ``predicate``."""

    def transform(upstream: Upstream[In], _scope: CloseableScope) -> AnyEffect:
        taken: list[In] = []
        leftover: Chunk | None = None

        def step(chunk: tuple[In, ...]) -> AnyEffect:
            def body() -> Generator[AnyEffect, Any, AnyEffect]:
                nonlocal leftover
                for i, item in enumerate(chunk):
                    matched = yield predicate(item)
                    taken.append(item)
                    if matched:
                        leftover = _rest(chunk, i + 1)
                        return ENDED
                return VOID

            return _Gen(body).flatten()

        return read_until_done(upstream.flat_map(step), lambda: (taken, leftover))

    return Sink(transform)

take_while(predicate) -> Sink[list[Any], Any, Any]

def take_while(predicate: Callable[[Any], bool]) -> Sink[list[Any], Any, Any]:
    """A sink taking items while ``predicate`` holds."""
    return _taking_while(lambda item: Ok(item) if predicate(item) else _REJECTED)

take_while_effect(predicate) -> Sink[list[In], In, In, E, R]

def take_while_effect[In, E = Never, R = Never](
    predicate: Callable[[In], Effect[bool, E, R]],
) -> Sink[list[In], In, In, E, R]:
    """``take_while`` with an effectful ``predicate``."""
    return _taking_while_effect(
        lambda item: predicate(item).map(lambda ok: Ok(item) if ok else _REJECTED)
    )

take_while_filter(f) -> Sink[list[B], In, In]

def take_while_filter[In, B](f: Callable[[In], Result[B, object]]) -> Sink[list[B], In, In]:
    """A sink taking the ``Ok`` values of ``f(item)`` until it gives an
    ``Err``."""
    return _taking_while(f)

take_while_filter_effect(f) -> Sink[list[B], In, In, E, R]

def take_while_filter_effect[In, B, E = Never, R = Never](
    f: Callable[[In], Effect[Result[B, object], E, R]],
) -> Sink[list[B], In, In, E, R]:
    """``take_while_filter`` with an effectful ``f``."""
    return _taking_while_effect(f)