add .flatten_ok()

Author: vadixidav

src/flatten_ok.rs

@@ -0,0 +1,61 @@
+pub fn flatten_ok<I, T, E>(iter: I) -> FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+    FlattenOk { iter, inner: None }
+}
+
+/// An iterator adaptor that flattens `Result::Ok` values and
+/// allows `Result::Err` values through unchanged.
+///
+/// See [`.flatten_ok()`](crate::Itertools::flatten_ok) for more information.
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+    iter: I,
+    inner: Option<T::IntoIter>,
+}
+
+impl<I, T, E> Iterator for FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+    type Item = Result<T::Item, E>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some(inner) = &mut self.inner {
+                if let Some(item) = inner.next() {
+                    return Some(Ok(item));
+                } else {
+                    self.inner = None;
+                }
+            } else {
+                match self.iter.next() {
+                    Some(Ok(ok)) => self.inner = Some(ok.into_iter()),
+                    Some(Err(e)) => return Some(Err(e)),
+                    None => return None,
+                }
+            }
+        }
+    }
+}
+
+impl<I, T, E> Clone for FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>> + Clone,
+    T: IntoIterator,
+    T::IntoIter: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            iter: self.iter.clone(),
+            inner: self.inner.clone(),
+        }
+    }
+}

src/lib.rs

@@ -119,6 +119,7 @@ pub mod structs {
     pub use crate::cons_tuples_impl::ConsTuples;
     pub use crate::exactly_one_err::ExactlyOneError;
     pub use crate::format::{Format, FormatWith};
+    pub use crate::flatten_ok::FlattenOk;
     #[cfg(feature = "use_std")]
     pub use crate::grouping_map::{GroupingMap, GroupingMapBy};
     #[cfg(feature = "use_alloc")]
@@ -194,6 +195,7 @@ mod combinations;
 mod combinations_with_replacement;
 mod exactly_one_err;
 mod diff;
+mod flatten_ok;
 mod format;
 #[cfg(feature = "use_std")]
 mod grouping_map;
@@ -833,6 +835,30 @@ pub trait Itertools : Iterator {
         adaptors::filter_map_ok(self, f)
     }
 
+    /// Return an iterator adaptor that flattens every `Result::Ok` value into
+    /// a series of `Result::Ok` values. `Result::Err` values are unchanged.
+    /// 
+    /// This is useful when you have some common error type for your crate and
+    /// need to propogate it upwards, but the `Result::Ok` case needs to be flattened.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let input = vec![Ok(0..2), Err(false), Ok(2..4)];
+    /// let it = input.iter().cloned().flatten_ok();
+    /// itertools::assert_equal(it.clone(), vec![Ok(0), Ok(1), Err(false), Ok(2), Ok(3)]);
+    /// 
+    /// // This can also be used to propogate errors when collecting.
+    /// let output_result: Result<Vec<i32>, bool> = it.collect();
+    /// assert_eq!(output_result, Err(false));
+    /// ```
+    fn flatten_ok<T, E>(self) -> FlattenOk<Self, T, E>
+        where Self: Iterator<Item = Result<T, E>> + Sized,
+              T: IntoIterator
+    {
+        flatten_ok::flatten_ok(self)
+    }
+
     /// Return an iterator adaptor that merges the two base iterators in
     /// ascending order.  If both base iterators are sorted (ascending), the
     /// result is sorted.