Merge #527

527: add .flatten_ok() r=jswrenn a=vadixidav

Closes #522

Co-authored-by: Geordon Worley <[email protected]>

Author: bors[bot]

src/flatten_ok.rs

@@ -0,0 +1,166 @@
+use crate::size_hint;
+use std::{
+    fmt,
+    iter::{DoubleEndedIterator, FusedIterator},
+};
+
+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_front: None,
+        inner_back: 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_front: Option<T::IntoIter>,
+    inner_back: 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 {
+            // Handle the front inner iterator.
+            if let Some(inner) = &mut self.inner_front {
+                if let Some(item) = inner.next() {
+                    return Some(Ok(item));
+                } else {
+                    // This is necessary for the iterator to implement `FusedIterator`
+                    // with only the orginal iterator being fused.
+                    self.inner_front = None;
+                }
+            }
+
+            match self.iter.next() {
+                Some(Ok(ok)) => self.inner_front = Some(ok.into_iter()),
+                Some(Err(e)) => return Some(Err(e)),
+                None => {
+                    // Handle the back inner iterator.
+                    if let Some(inner) = &mut self.inner_back {
+                        if let Some(item) = inner.next() {
+                            return Some(Ok(item));
+                        } else {
+                            // This is necessary for the iterator to implement `FusedIterator`
+                            // with only the orginal iterator being fused.
+                            self.inner_back = None;
+                        }
+                    } else {
+                        return None;
+                    }
+                }
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let inner_hint = |inner: &Option<T::IntoIter>| {
+            inner
+                .as_ref()
+                .map(Iterator::size_hint)
+                .unwrap_or((0, Some(0)))
+        };
+        let inner_front = inner_hint(&self.inner_front);
+        let inner_back = inner_hint(&self.inner_back);
+        // The outer iterator `Ok` case could be (0, None) as we don't know its size_hint yet.
+        let outer = match self.iter.size_hint() {
+            (0, Some(0)) => (0, Some(0)),
+            _ => (0, None),
+        };
+
+        size_hint::add(size_hint::add(inner_front, inner_back), outer)
+    }
+}
+
+impl<I, T, E> DoubleEndedIterator for FlattenOk<I, T, E>
+where
+    I: DoubleEndedIterator<Item = Result<T, E>>,
+    T: IntoIterator,
+    T::IntoIter: DoubleEndedIterator,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        loop {
+            // Handle the back inner iterator.
+            if let Some(inner) = &mut self.inner_back {
+                if let Some(item) = inner.next_back() {
+                    return Some(Ok(item));
+                } else {
+                    // This is necessary for the iterator to implement `FusedIterator`
+                    // with only the orginal iterator being fused.
+                    self.inner_back = None;
+                }
+            }
+
+            match self.iter.next_back() {
+                Some(Ok(ok)) => self.inner_back = Some(ok.into_iter()),
+                Some(Err(e)) => return Some(Err(e)),
+                None => {
+                    // Handle the front inner iterator.
+                    if let Some(inner) = &mut self.inner_front {
+                        if let Some(item) = inner.next_back() {
+                            return Some(Ok(item));
+                        } else {
+                            // This is necessary for the iterator to implement `FusedIterator`
+                            // with only the orginal iterator being fused.
+                            self.inner_front = None;
+                        }
+                    } else {
+                        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,
+{
+    #[inline]
+    clone_fields!(iter, inner_front, inner_back);
+}
+
+impl<I, T, E> fmt::Debug for FlattenOk<I, T, E>
+where
+    I: Iterator<Item = Result<T, E>> + fmt::Debug,
+    T: IntoIterator,
+    T::IntoIter: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("FlattenOk")
+            .field("iter", &self.iter)
+            .field("inner_front", &self.inner_front)
+            .field("inner_back", &self.inner_back)
+            .finish()
+    }
+}
+
+/// Only the iterator being flattened needs to implement [`FusedIterator`].
+impl<I, T, E> FusedIterator for FlattenOk<I, T, E>
+where
+    I: FusedIterator<Item = Result<T, E>>,
+    T: IntoIterator,
+{
+}

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;
@@ -899,6 +901,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.

tests/flatten_ok.rs

@@ -0,0 +1,76 @@
+use itertools::{assert_equal, Itertools};
+use std::{ops::Range, vec::IntoIter};
+
+fn mix_data() -> IntoIter<Result<Range<i32>, bool>> {
+    vec![Ok(0..2), Err(false), Ok(2..4), Err(true), Ok(4..6)].into_iter()
+}
+
+fn ok_data() -> IntoIter<Result<Range<i32>, bool>> {
+    vec![Ok(0..2), Ok(2..4), Ok(4..6)].into_iter()
+}
+
+#[test]
+fn flatten_ok_mixed_expected_forward() {
+    assert_equal(
+        mix_data().flatten_ok(),
+        vec![
+            Ok(0),
+            Ok(1),
+            Err(false),
+            Ok(2),
+            Ok(3),
+            Err(true),
+            Ok(4),
+            Ok(5),
+        ],
+    );
+}
+
+#[test]
+fn flatten_ok_mixed_expected_reverse() {
+    assert_equal(
+        mix_data().flatten_ok().rev(),
+        vec![
+            Ok(5),
+            Ok(4),
+            Err(true),
+            Ok(3),
+            Ok(2),
+            Err(false),
+            Ok(1),
+            Ok(0),
+        ],
+    );
+}
+
+#[test]
+fn flatten_ok_collect_mixed_forward() {
+    assert_eq!(
+        mix_data().flatten_ok().collect::<Result<Vec<_>, _>>(),
+        Err(false)
+    );
+}
+
+#[test]
+fn flatten_ok_collect_mixed_reverse() {
+    assert_eq!(
+        mix_data().flatten_ok().rev().collect::<Result<Vec<_>, _>>(),
+        Err(true)
+    );
+}
+
+#[test]
+fn flatten_ok_collect_ok_forward() {
+    assert_eq!(
+        ok_data().flatten_ok().collect::<Result<Vec<_>, _>>(),
+        Ok((0..6).collect())
+    );
+}
+
+#[test]
+fn flatten_ok_collect_ok_reverse() {
+    assert_eq!(
+        ok_data().flatten_ok().rev().collect::<Result<Vec<_>, _>>(),
+        Ok((0..6).rev().collect())
+    );
+}