AsyncHandleEventsSequence.swift

///  .handleEvents(onNext: { print($0) }) Similar to the Combine variant, but for AsyncSequence.
///

public extension AsyncSequence {

    /// Attach event handlers to interation events of the sequence.
    /// - Parameters:
    ///   - onInit: handler called when interator is created
    ///   - onNext: handler called when value is received
    ///   - onCompletion: handler called when sequence competes
    func handleEvents(onInit: (@Sendable () -> Void)? = nil,
                      onNext: (@Sendable (Element) -> Void)? = nil,
                      onCompletion: (@Sendable (AsyncHandleEventsSequence<Self>.Completion) -> Void)? = nil) -> AsyncHandleEventsSequence<Self> {
        AsyncHandleEventsSequence(
            upstream: self,
            onInit: onInit,
            onNext: onNext,
            onCompletion: onCompletion
        )
    }
}

public struct AsyncHandleEventsSequence<Upstream>: AsyncSequence where Upstream: AsyncSequence {
    public typealias Element = Upstream.Element

    private let upstream: Upstream
    private let onInit: (@Sendable () -> Void)?
    private let onNext: (@Sendable (Element) -> Void)?
    private let onCompletion: (@Sendable (Completion) -> Void)?

    public enum Completion {
        case finished
        case failure(Error)
    }

    /// Creates a sequence that attaches handlers to all interation events.
    /// - Parameters:
    ///   - onInit: handler called when interator is created
    ///   - onNext: handler called when value is received
    ///   - onCompletion: handler called when sequence competes
    public init(upstream: Upstream,
                onInit: (@Sendable () -> Void)? = nil,
                onNext: (@Sendable (Element) -> Void)? = nil,
                onCompletion: (@Sendable (Completion) -> Void)? = nil) {
        self.upstream = upstream
        self.onInit = onInit
        self.onNext = onNext
        self.onCompletion = onCompletion
    }

    public func makeAsyncIterator() -> Iterator<Upstream.AsyncIterator> {
        Iterator(
            upstream.makeAsyncIterator(),
            onInit: onInit,
            onNext: onNext,
            onCompletion: onCompletion
        )
    }

    public struct Iterator<Inner>: AsyncIteratorProtocol where Inner: AsyncIteratorProtocol, Element == Inner.Element {

        private var it: Inner
        private let onNext: (@Sendable (Element) -> Void)?
        private let onCompletion: (@Sendable (Completion) -> Void)?

        fileprivate init(_ it: Inner,
                         onInit: (@Sendable () -> Void)?,
                         onNext: (@Sendable (Element) -> Void)?,
                         onCompletion: (@Sendable (Completion) -> Void)?) {
            self.it = it
            onInit?()
            self.onNext = onNext
            self.onCompletion = onCompletion
        }

        public mutating func next() async throws -> Inner.Element? {
            do {
                let element = try await it.next()
                if let element = element {
                    onNext?(element)
                } else {
                    onCompletion?(.finished)
                }
                return element
            } catch {
                onCompletion?(.failure(error))
                throw error
            }
        }
    }
}