try_stream_transpose/
lib.rs

1#![doc = include_str!("../README.md")]
2use core::fmt::{self, Debug, Display};
3use futures_util::{FutureExt, Stream, StreamExt, TryStream, future};
4use tokio::{
5    sync::mpsc::{self, error::SendError},
6    task::JoinError,
7};
8use tokio_stream::wrappers::ReceiverStream;
9
10/// Convenience for `TryStream`, when you want to work on a stream of `Ok`s as
11/// a stream of the values inside the `Ok`s.
12pub trait TryStreamTranspose: TryStream {
13    /// Transposes a [`TryStream`], in the sense that it passes a stream of its
14    /// unwrapped items to the closure `f`.
15    ///
16    /// If, while executing the closure, one of the items turns out an `Err`
17    /// value, the in-closure stream stops yielding and this function returns
18    /// said error. If neither the stream nor this function, while executing,
19    /// produce an error, the closure's result is returned.
20    ///
21    /// Error types returned by this function must implement conversions from
22    /// [`Error`], in addition to the conversion from the items's error type.
23    /// Additionally, the `buf_size` parameter controls the size of the internal
24    /// buffer.
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 `Error`. Conversion from
52    /// // `&'static str` is implemented to support our code above.
53    /// #[derive(Debug, thiserror::Error)]
54    /// #[error(transparent)]
55    /// enum Error {
56    ///     #[error("{0}")]
57    ///     Message(&'static str),
58    ///     Internal(#[from] try_stream_transpose::Error<i32>),
59    /// }
60    ///
61    /// impl From<&'static str> for Error {
62    ///     fn from(msg: &'static str) -> Self {
63    ///         Self::Message(msg)
64    ///     }
65    /// }
66    /// # }
67    /// # tokio::runtime::Runtime::new().unwrap().block_on(asd())
68    /// ```
69    fn transpose_with<Fut, F, T, E>(
70        self,
71        buf_size: usize,
72        mut f: F,
73    ) -> impl Future<Output = Result<T, E>>
74    where
75        F: FnMut(ReceiverStream<Self::Ok>) -> Fut,
76        Fut: Future<Output = Result<T, E>>,
77        Self: Send
78            + Sized
79            + Stream<Item = Result<Self::Ok, Self::Error>>
80            + 'static,
81        Self::Ok: Send,
82        Self::Error: Send,
83        E: Send + From<Self::Error> + From<Error<Self::Ok>> + 'static,
84    {
85        // To unwrap the items from their `Result`s and not evaluate the entire
86        // stream at once, we're sending the items through a channel and
87        // returning errors early once we're encountering them.
88        let (sender, recver) = mpsc::channel(buf_size);
89
90        // Create a separate task where we're sending the stream's items through
91        // the channel and either return errors, or `()` on success.
92        let send_handle = tokio::spawn(async move {
93            tokio::pin! {
94                let stream = self;
95            }
96            while let Some(line) = stream.next().await {
97                sender.send(line?).await.map_err(Error::from)?;
98            }
99            Ok(())
100        });
101        // Wrap the channel's receiver in a `ReceiverStream` and pass that to
102        // the closure.
103        let recv_handle = f(ReceiverStream::new(recver));
104
105        // Join the futures for the sender and the receiver task.
106        future::join(send_handle, recv_handle).map(|result| match result {
107            // Joining the sender failed.
108            (Err(err), _) => Err(Error::from(err).into()),
109            // Joining the sender succeeded, but sending failed somewhere or the
110            // stream contained an `Err` value.
111            (Ok(Err(err)), _) => Err(err),
112            // Return the computation result.
113            (_, result) => result,
114        })
115    }
116}
117
118/// Abstract the errors encountered internally.
119#[derive(Debug)]
120pub struct Error<Item>(ErrorRepr<Item>);
121
122#[derive(Debug)]
123pub(crate) enum ErrorRepr<Item> {
124    Join(JoinError),
125    Send(SendError<Item>),
126}
127
128impl<Item> Display for Error<Item> {
129    #[inline]
130    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
131        Display::fmt("unable to transpose item", f)
132    }
133}
134
135impl<Item> core::error::Error for Error<Item>
136where
137    Item: Debug + 'static,
138{
139    #[inline]
140    fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
141        match &self.0 {
142            ErrorRepr::Join(err) => Some(err),
143            ErrorRepr::Send(err) => Some(err),
144        }
145    }
146}
147
148impl<S: TryStream> TryStreamTranspose for S {}
149
150impl<Item> From<JoinError> for Error<Item> {
151    #[inline]
152    fn from(value: JoinError) -> Self {
153        Self(ErrorRepr::Join(value))
154    }
155}
156
157impl<Item> From<SendError<Item>> for Error<Item> {
158    #[inline]
159    fn from(value: SendError<Item>) -> Self {
160        Self(ErrorRepr::Send(value))
161    }
162}
try_stream_transpose/lib.rs

try_stream_transpose/
lib.rs
