7971: stream: add StreamExt::filter_map_async

tokio-stream/src/lib.rs

@@ -80,8 +80,8 @@ pub use stream_ext::{collect::FromStream, StreamExt};
 /// Adapters for [`Stream`]s created by methods in [`StreamExt`].
 pub mod adapters {
     pub use crate::stream_ext::{
-        Chain, Filter, FilterMap, Fuse, Map, MapWhile, Merge, Peekable, Skip, SkipWhile, Take,
-        TakeWhile, Then,
+        Chain, Filter, FilterMap, FilterMapAsync, Fuse, Map, MapWhile, Merge, Peekable, Skip,
+        SkipWhile, Take, TakeWhile, Then,
     };
     cfg_time! {
         pub use crate::stream_ext::{ChunksTimeout, Timeout, TimeoutRepeating};

tokio-stream/src/stream_ext.rs

@@ -16,6 +16,9 @@ use collect::{Collect, FromStream};
 mod filter;
 pub use filter::Filter;
 
+mod filter_map_async;
+pub use filter_map_async::FilterMapAsync;
+
 mod filter_map;
 pub use filter_map::FilterMap;
 
@@ -479,6 +482,53 @@ pub trait StreamExt: Stream {
         FilterMap::new(self, f)
     }
 
+    /// Filters the values produced by this stream asynchronously while
+    /// simultaneously mapping them to a different type according to the
+    /// provided async closure.
+    ///
+    /// The provided closure is executed over all elements of this stream as
+    /// they are made available, and the returned future is executed. Only one
+    /// future is executed at the time. If the returned future resolves to
+    /// [`Some(item)`](Some) then the stream will yield the value `item`, but if
+    /// it resolves to [`None`], then the value will be skipped.
+    ///
+    /// Note that this function consumes the stream passed into it and returns a
+    /// wrapped version of it, similar to [`Iterator::filter_map`] method in the
+    /// standard library.
+    ///
+    /// Be aware that if the future is not `Unpin`, then neither is the `Stream`
+    /// returned by this method. To handle this, you can use `tokio::pin!` as in
+    /// the example below or put the stream in a `Box` with `Box::pin(stream)`.
+    ///
+    /// # Examples
+    /// ```
+    /// # #[tokio::main(flavor = "current_thread")]
+    /// # async fn main() {
+    /// use tokio_stream::{self as stream, StreamExt};
+    ///
+    /// let stream = stream::iter(0..=7);
+    /// let odds = stream.filter_map_async(async |x| {
+    ///     if x % 2 == 0 { Some(x + 1) } else { None }
+    /// });
+    ///
+    /// tokio::pin!(odds);
+    ///
+    /// assert_eq!(Some(1), odds.next().await);
+    /// assert_eq!(Some(3), odds.next().await);
+    /// assert_eq!(Some(5), odds.next().await);
+    /// assert_eq!(Some(7), odds.next().await);
+    /// assert_eq!(None, odds.next().await);
+    /// # }
+    /// ```
+    fn filter_map_async<T, F, Fut>(self, f: F) -> FilterMapAsync<Self, Fut, F>
+    where
+        F: FnMut(Self::Item) -> Fut,
+        Fut: Future<Output = Option<T>>,
+        Self: Sized,
+    {
+        FilterMapAsync::new(self, f)
+    }
+
     /// Creates a stream which ends after the first `None`.
     ///
     /// After a stream returns `None`, behavior is undefined. Future calls to

tokio-stream/src/stream_ext/filter_map_async.rs

@@ -0,0 +1,75 @@
+use crate::Stream;
+
+use core::fmt;
+use core::future::Future;
+use core::pin::Pin;
+use core::task::{ready, Context, Poll};
+use pin_project_lite::pin_project;
+
+pin_project! {
+    /// Stream for the [`filter_map_async`](super::StreamExt::filter_map_async) method.
+    #[must_use = "streams do nothing unless polled"]
+    pub struct FilterMapAsync<St, Fut, F> {
+        #[pin]
+        stream: St,
+        #[pin]
+        future: Option<Fut>,
+        f: F,
+    }
+}
+
+impl<St, Fut, F> fmt::Debug for FilterMapAsync<St, Fut, F>
+where
+    St: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("FilterMapAsync")
+            .field("stream", &self.stream)
+            .finish()
+    }
+}
+
+impl<St, Fut, F> FilterMapAsync<St, Fut, F> {
+    pub(super) fn new(stream: St, f: F) -> Self {
+        FilterMapAsync {
+            stream,
+            future: None,
+            f,
+        }
+    }
+}
+
+impl<T, St, F, Fut> Stream for FilterMapAsync<St, Fut, F>
+where
+    St: Stream,
+    Fut: Future<Output = Option<T>>,
+    F: FnMut(St::Item) -> Fut,
+{
+    type Item = T;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
+        let mut me = self.project();
+
+        loop {
+            if let Some(future) = me.future.as_mut().as_pin_mut() {
+                if let Some(item) = ready!(future.poll(cx)) {
+                    me.future.set(None);
+                    return Poll::Ready(Some(item));
+                } else {
+                    me.future.set(None);
+                }
+            }
+
+            match ready!(me.stream.as_mut().poll_next(cx)) {
+                Some(item) => {
+                    me.future.set(Some((me.f)(item)));
+                }
+                None => return Poll::Ready(None),
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (0, self.stream.size_hint().1)
+    }
+}