try_stream_transpose/

lib.rs

#![doc = include_str!("../README.md")]
use core::fmt::{self, Debug, Display};
use futures_util::{FutureExt, Stream, StreamExt, TryStream, future};
use tokio::{
    sync::mpsc::{self, error::SendError},
    task::JoinError,
};
use tokio_stream::wrappers::ReceiverStream;

/// Convenience for `TryStream`, when you want to work on a stream of `Ok`s as
/// a stream of the values inside the `Ok`s.
pub trait TryStreamTranspose: TryStream {
    /// Transposes a [`TryStream`], in the sense that it passes a stream of its
    /// unwrapped items to the closure `f`.
    ///
    /// If, while executing the closure, one of the items turns out an `Err`
    /// value, the in-closure stream stops yielding and this function returns
    /// said error. If neither the stream nor this function, while executing,
    /// produce an error, the closure's result is returned.
    ///
    /// Error types returned by this function must implement conversions from
    /// [`Error`], in addition to the conversion from the items's error type.
    /// Additionally, the `buf_size` parameter controls the size of the internal
    /// buffer.
    /// ```
    /// # async fn asd() {
    /// use futures_util::{FutureExt, StreamExt, TryStreamExt, stream};
    /// use try_stream_transpose::TryStreamTranspose;
    ///
    /// let nums: [Result<i32, Error>; 4] = [Ok(4), Ok(5), Ok(6), Ok(7)];
    /// core::assert_matches!(
    ///     stream::iter(nums)
    ///         .map_ok(|x| x + 1)
    ///         .transpose_with(1024, |s| {
    ///             s.fold(1, async |a, b| a * b).map(|x| Ok(x / 40))
    ///         })
    ///         .await,
    ///     Ok::<i32, Error>(42)
    /// );
    /// let argh = [Ok(2), Err("3"), Ok(7)];
    /// core::assert_matches!(
    ///     stream::iter(argh)
    ///         .map_ok(|x| x + 1)
    ///         .transpose_with(1024, |s| {
    ///             s.fold(1, async |a, b| a * b).map(|x| Ok(x / 8))
    ///         })
    ///         .await,
    ///     Err(Error::Message("3"))
    /// );
    ///
    /// // The error type must support conversion from `Error`. Conversion from
    /// // `&'static str` is implemented to support our code above.
    /// #[derive(Debug, thiserror::Error)]
    /// #[error(transparent)]
    /// enum Error {
    ///     #[error("{0}")]
    ///     Message(&'static str),
    ///     Internal(#[from] try_stream_transpose::Error<i32>),
    /// }
    ///
    /// impl From<&'static str> for Error {
    ///     fn from(msg: &'static str) -> Self {
    ///         Self::Message(msg)
    ///     }
    /// }
    /// # }
    /// # tokio::runtime::Runtime::new().unwrap().block_on(asd())
    /// ```
    fn transpose_with<Fut, F, T, E>(
        self,
        buf_size: usize,
        mut f: F,
    ) -> impl Future<Output = Result<T, E>>
    where
        F: FnMut(ReceiverStream<Self::Ok>) -> Fut,
        Fut: Future<Output = Result<T, E>>,
        Self: Send
            + Sized
            + Stream<Item = Result<Self::Ok, Self::Error>>
            + 'static,
        Self::Ok: Send,
        Self::Error: Send,
        E: Send + From<Self::Error> + From<Error<Self::Ok>> + 'static,
    {
        // To unwrap the items from their `Result`s and not evaluate the entire
        // stream at once, we're sending the items through a channel and
        // returning errors early once we're encountering them.
        let (sender, recver) = mpsc::channel(buf_size);

        // Create a separate task where we're sending the stream's items through
        // the channel and either return errors, or `()` on success.
        let send_handle = tokio::spawn(async move {
            tokio::pin! {
                let stream = self;
            }
            while let Some(line) = stream.next().await {
                sender.send(line?).await.map_err(InternalSendError::Send)?;
            }
            Ok(())
        });
        // Wrap the channel's receiver in a `ReceiverStream` and pass that to
        // the closure.
        let recv_handle = f(ReceiverStream::new(recver));

        // Join the futures for the sender and the receiver task.
        future::join(send_handle, recv_handle).map(|result| match result {
            // Joining the sender failed.
            (Err(err), _) => Err(Error::from(err).into()),
            // Joining the sender succeeded, but a stream item was an error.
            (Ok(Err(InternalSendError::Item(err))), _) => Err(err.into()),
            // Processing the stream returned an error.
            (_, Err(err)) => Err(err),
            // Joining the sender succeeded, but sending failed somewhere.
            (Ok(Err(InternalSendError::Send(err))), _) => {
                Err(Error::from(err).into())
            }
            // Return the successful computation result.
            (_, ok) => ok,
        })
    }
}

/// Abstract the errors encountered internally.
#[derive(Debug)]
pub struct Error<Item>(ErrorRepr<Item>);

#[derive(Debug)]
pub(crate) enum ErrorRepr<Item> {
    Join(JoinError),
    Send(SendError<Item>),
}

enum InternalSendError<I, E> {
    Item(E),
    Send(SendError<I>),
}

impl<Item> Display for Error<Item> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        Display::fmt("unable to transpose item", f)
    }
}

impl<Item> core::error::Error for Error<Item>
where
    Item: Debug + 'static,
{
    #[inline]
    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
        match &self.0 {
            ErrorRepr::Join(err) => Some(err),
            ErrorRepr::Send(err) => Some(err),
        }
    }
}

impl<S: TryStream> TryStreamTranspose for S {}

impl<Item> From<JoinError> for Error<Item> {
    #[inline]
    fn from(value: JoinError) -> Self {
        Self(ErrorRepr::Join(value))
    }
}

impl<Item> From<SendError<Item>> for Error<Item> {
    #[inline]
    fn from(value: SendError<Item>) -> Self {
        Self(ErrorRepr::Send(value))
    }
}

impl<I, E> From<E> for InternalSendError<I, E> {
    #[inline]
    fn from(err: E) -> Self {
        Self::Item(err)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures_util::stream;

    #[tokio::test]
    async fn error_propagation() {
        #[derive(Debug, thiserror::Error)]
        #[error(transparent)]
        enum TestError {
            #[error("{0}")]
            Message(&'static str),
            Internal(#[from] Error<i32>),
        }
        let nums: [Result<i32, TestError>; 4] = [Ok(4), Ok(5), Ok(6), Ok(7)];
        core::assert_matches!(
            stream::iter(nums)
                .transpose_with(1024, |_| {
                    future::err(TestError::Message("stream err"))
                })
                .await,
            Err::<i32, TestError>(TestError::Message("stream err")),
            "error propagation from inside the closure"
        );
    }
}