Merge pull request #1141 from cramertj/map-ok

Add TryStream::{map_ok, map_err}

Author: MajorBreakfast

futures-util/src/try_stream/map_err.rs

@@ -0,0 +1,44 @@
+use core::marker::Unpin;
+use core::mem::PinMut;
+use futures_core::stream::{Stream, TryStream};
+use futures_core::task::{self, Poll};
+
+/// Stream for the [`map_err`](super::TryStreamExt::map_err) combinator.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct MapErr<St, F> {
+    stream: St,
+    f: F,
+}
+
+impl<St, F> MapErr<St, F> {
+    unsafe_pinned!(stream: St);
+    unsafe_unpinned!(f: F);
+
+    /// Creates a new MapErr.
+    pub(super) fn new(stream: St, f: F) -> Self {
+        MapErr { stream, f }
+    }
+}
+
+impl<St: Unpin, F> Unpin for MapErr<St, F> {}
+
+impl<St, F, E> Stream for MapErr<St, F>
+where
+    St: TryStream,
+    F: FnMut(St::Error) -> E,
+{
+    type Item = Result<St::Ok, E>;
+
+    #[allow(redundant_closure)] // https://github.com/rust-lang-nursery/rust-clippy/issues/1439
+    fn poll_next(
+        mut self: PinMut<Self>,
+        cx: &mut task::Context,
+    ) -> Poll<Option<Self::Item>> {
+        match self.stream().try_poll_next(cx) {
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(opt) =>
+                Poll::Ready(opt.map(|res| res.map_err(|e| self.f()(e)))),
+        }
+    }
+}

futures-util/src/try_stream/map_ok.rs

@@ -0,0 +1,44 @@
+use core::marker::Unpin;
+use core::mem::PinMut;
+use futures_core::stream::{Stream, TryStream};
+use futures_core::task::{self, Poll};
+
+/// Stream for the [`map_ok`](super::TryStreamExt::map_ok) combinator.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct MapOk<St, F> {
+    stream: St,
+    f: F,
+}
+
+impl<St, F> MapOk<St, F> {
+    unsafe_pinned!(stream: St);
+    unsafe_unpinned!(f: F);
+
+    /// Creates a new MapOk.
+    pub(super) fn new(stream: St, f: F) -> Self {
+        MapOk { stream, f }
+    }
+}
+
+impl<St: Unpin, F> Unpin for MapOk<St, F> {}
+
+impl<St, F, T> Stream for MapOk<St, F>
+where
+    St: TryStream,
+    F: FnMut(St::Ok) -> T,
+{
+    type Item = Result<T, St::Error>;
+
+    #[allow(redundant_closure)] // https://github.com/rust-lang-nursery/rust-clippy/issues/1439
+    fn poll_next(
+        mut self: PinMut<Self>,
+        cx: &mut task::Context,
+    ) -> Poll<Option<Self::Item>> {
+        match self.stream().try_poll_next(cx) {
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(opt) =>
+                Poll::Ready(opt.map(|res| res.map(|x| self.f()(x)))),
+        }
+    }
+}

futures-util/src/try_stream/mod.rs

@@ -10,6 +10,12 @@ use futures_core::stream::TryStream;
 mod err_into;
 pub use self::err_into::ErrInto;
 
+mod map_ok;
+pub use self::map_ok::MapOk;
+
+mod map_err;
+pub use self::map_err::MapErr;
+
 mod try_next;
 pub use self::try_next::TryNext;
 
@@ -51,6 +57,58 @@ pub trait TryStreamExt: TryStream {
         ErrInto::new(self)
     }
 
+    /// Wraps the current stream in a new stream which maps the success value
+    /// using the provided closure.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await, await_macro)]
+    /// # futures::executor::block_on(async {
+    /// use futures::{stream, TryStreamExt};
+    ///
+    /// let mut stream =
+    ///     stream::iter(vec![Ok(5), Err(0)])
+    ///         .map_ok(|x| x + 2);
+    ///
+    /// assert_eq!(await!(stream.try_next()), Ok(Some(7)));
+    /// assert_eq!(await!(stream.try_next()), Err(0));
+    /// # })
+    /// ```
+    fn map_ok<T, F>(self, f: F) -> MapOk<Self, F>
+    where
+        Self: Sized,
+        F: FnMut(Self::Ok) -> T,
+    {
+        MapOk::new(self, f)
+    }
+
+    /// Wraps the current stream in a new stream which maps the error value
+    /// using the provided closure.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await, await_macro)]
+    /// # futures::executor::block_on(async {
+    /// use futures::{stream, TryStreamExt};
+    ///
+    /// let mut stream =
+    ///     stream::iter(vec![Ok(5), Err(0)])
+    ///         .map_err(|x| x + 2);
+    ///
+    /// assert_eq!(await!(stream.try_next()), Ok(Some(5)));
+    /// assert_eq!(await!(stream.try_next()), Err(2));
+    /// # })
+    /// ```
+    fn map_err<E, F>(self, f: F) -> MapErr<Self, F>
+    where
+        Self: Sized,
+        F: FnMut(Self::Error) -> E,
+    {
+        MapErr::new(self, f)
+    }
+
     /// Creates a future that attempts to resolve the next item in the stream.
     /// If an error is encountered before the next item, the error is returned
     /// instead.