try_stream_transpose/
lib.rs

1#![doc = include_str!("../README.md")]
2use futures_util::{FutureExt, Stream, StreamExt, TryStream, future};
3use tokio::{
4    sync::mpsc::{self, error::SendError},
5    task::JoinError,
6};
7use tokio_stream::wrappers::ReceiverStream;
8
9/// Convenience for `TryStream`, when you want to work on a stream of `Ok`s as
10/// a stream of the values inside the `Ok`s.
11pub trait TryStreamTranspose: TryStream {
12    /// Transposes a [`TryStream`], in the sense that it passes a stream of its
13    /// unwrapped items to the closure `f`.
14    ///
15    /// If, while executing the closure, one of the items turns out an `Err`
16    /// value, the in-closure stream stops yielding and this function returns
17    /// said error. If neither the stream nor this function, while executing,
18    /// produce an error, the closure's result is returned.
19    ///
20    /// Due to implementation details, errors returned by this function --- it
21    /// uses a channel internally --- must implement conversions from
22    /// [`JoinError`] and [`SendError`], in addition to the conversion from the
23    /// items's error type. Additionally, the `buf_size` parameter controls the
24    /// size of the channel's queue.
25    /// ```
26    /// # async fn asd() {
27    /// use futures_util::{FutureExt, StreamExt, TryStreamExt, stream};
28    /// use try_stream_transpose::TryStreamTranspose;
29    ///
30    /// let nums: [Result<i32, Error>; 4] = [Ok(4), Ok(5), Ok(6), Ok(7)];
31    /// core::assert_matches!(
32    ///     stream::iter(nums)
33    ///         .map_ok(|x| x + 1)
34    ///         .transpose_with(1024, |s| {
35    ///             s.fold(1, async |a, b| a * b).map(|x| Ok(x / 40))
36    ///         })
37    ///         .await,
38    ///     Ok::<i32, Error>(42)
39    /// );
40    /// let argh = [Ok(2), Err("3"), Ok(7)];
41    /// core::assert_matches!(
42    ///     stream::iter(argh)
43    ///         .map_ok(|x| x + 1)
44    ///         .transpose_with(1024, |s| {
45    ///             s.fold(1, async |a, b| a * b).map(|x| Ok(x / 8))
46    ///         })
47    ///         .await,
48    ///     Err(Error::Message("3"))
49    /// );
50    ///
51    /// // The error type must support conversion from `JoinError` and
52    /// // `SendError`. Conversion from `&'static str` is implemented to support
53    /// // our code above.
54    /// #[derive(Debug, thiserror::Error)]
55    /// #[error(transparent)]
56    /// enum Error {
57    ///     #[error("{0}")]
58    ///     Message(&'static str),
59    ///     JoinError(#[from] tokio::task::JoinError),
60    ///     SendError(#[from] tokio::sync::mpsc::error::SendError<i32>),
61    /// }
62    ///
63    /// impl From<&'static str> for Error {
64    ///     fn from(msg: &'static str) -> Self {
65    ///         Self::Message(msg)
66    ///     }
67    /// }
68    /// # }
69    /// # tokio::runtime::Runtime::new().unwrap().block_on(asd())
70    /// ```
71    fn transpose_with<Fut, F, T, E>(
72        self,
73        buf_size: usize,
74        mut f: F,
75    ) -> impl Future<Output = Result<T, E>>
76    where
77        F: FnMut(ReceiverStream<Self::Ok>) -> Fut,
78        Fut: Future<Output = Result<T, E>>,
79        Self: Send
80            + Sized
81            + Stream<Item = Result<Self::Ok, Self::Error>>
82            + 'static,
83        Self::Ok: Send,
84        Self::Error: Send,
85        E: Send
86            + From<Self::Error>
87            + From<JoinError>
88            + From<SendError<Self::Ok>>
89            + 'static,
90    {
91        // To unwrap the items from their `Result`s and not evaluate the entire
92        // stream at once, we're sending the items through a channel and
93        // returning errors early once we're encountering them.
94        let (sender, recver) = mpsc::channel(buf_size);
95
96        // Create a separate task where we're sending the stream's items through
97        // the channel and either return errors, or `()` on success.
98        let send_handle = tokio::spawn(async move {
99            tokio::pin! {
100                let stream = self;
101            }
102            while let Some(line) = stream.next().await {
103                sender.send(line?).await?;
104            }
105            Ok(())
106        });
107        // Wrap the channel's receiver in a `ReceiverStream` and pass that to
108        // the closure.
109        let recv_handle = f(ReceiverStream::new(recver));
110
111        // Join the futures for the sender and the receiver task.
112        future::join(send_handle, recv_handle).map(|result| match result {
113            // Joining the sender failed.
114            (Err(err), _) => Err(err.into()),
115            // Joining the sender succeeded, but sending failed somewhere or the
116            // stream contained an `Err` value.
117            (Ok(Err(err)), _) => Err(err),
118            // Return the computation result.
119            (_, result) => result,
120        })
121    }
122}
123
124impl<S: TryStream> TryStreamTranspose for S {}
try_stream_transpose/lib.rs

try_stream_transpose/
lib.rs
