try_stream-transpose 0.1.0

#![doc = include_str!("../README.md")]
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.
    ///
    /// Due to implementation details, errors returned by this function --- it
    /// uses a channel internally --- must implement conversions from
    /// [`JoinError`] and [`SendError`], in addition to the conversion from the
    /// items's error type. Additionally, the `buf_size` parameter controls the
    /// size of the channel's queue.
    /// ```
    /// # 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 `JoinError` and
    /// // `SendError`. Conversion from `&'static str` is implemented to support
    /// // our code above.
    /// #[derive(Debug, thiserror::Error)]
    /// #[error(transparent)]
    /// enum Error {
    ///     #[error("{0}")]
    ///     Message(&'static str),
    ///     JoinError(#[from] tokio::task::JoinError),
    ///     SendError(#[from] tokio::sync::mpsc::error::SendError<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<JoinError>
            + From<SendError<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?;
            }
            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(err.into()),
            // Joining the sender succeeded, but sending failed somewhere or the
            // stream contained an `Err` value.
            (Ok(Err(err)), _) => Err(err),
            // Return the computation result.
            (_, result) => result,
        })
    }
}

impl<S: TryStream> TryStreamTranspose for S {}