Merge #310

310: Add exactly_one function r=bluss a=Xaeroxe

Adds a small function I found uses for personally.  Can't be added to the core library because in the error case it allocates a vec and stores data in it.

Co-authored-by: Jacob Kiesel <[email protected]>

Author: bors[bot]

src/exactly_one_err.rs

@@ -0,0 +1,58 @@
+use std::iter::ExactSizeIterator;
+
+use size_hint;
+
+/// Iterator returned for the error case of `IterTools::exactly_one()`
+/// This iterator yields exactly the same elements as the input iterator.
+///
+/// During the execution of exactly_one the iterator must be mutated.  This wrapper
+/// effectively "restores" the state of the input iterator when it's handed back.
+///
+/// This is very similar to PutBackN except this iterator only supports 0-2 elements and does not
+/// use a `Vec`.
+#[derive(Debug, Clone)]
+pub struct ExactlyOneError<I>
+where
+    I: Iterator,
+{
+    first_two: (Option<I::Item>, Option<I::Item>),
+    inner: I,
+}
+
+impl<I> ExactlyOneError<I>
+where
+    I: Iterator,
+{
+    /// Creates a new `ExactlyOneErr` iterator.
+    pub fn new(first_two: (Option<I::Item>, Option<I::Item>), inner: I) -> Self {
+        Self { first_two, inner }
+    }
+}
+
+impl<I> Iterator for ExactlyOneError<I>
+where
+    I: Iterator,
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.first_two
+            .0
+            .take()
+            .or_else(|| self.first_two.1.take())
+            .or_else(|| self.inner.next())
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let mut additional_len = 0;
+        if self.first_two.0.is_some() {
+            additional_len += 1;
+        }
+        if self.first_two.1.is_some() {
+            additional_len += 1;
+        }
+        size_hint::add_scalar(self.inner.size_hint(), additional_len)
+    }
+}
+
+impl<I> ExactSizeIterator for ExactlyOneError<I> where I: ExactSizeIterator {}

src/lib.rs

@@ -102,6 +102,7 @@ pub mod structs {
     #[cfg(feature = "use_std")]
     pub use combinations::Combinations;
     pub use cons_tuples_impl::ConsTuples;
+    pub use exactly_one_err::ExactlyOneError;
     pub use format::{Format, FormatWith};
     #[cfg(feature = "use_std")]
     pub use groupbylazy::{IntoChunks, Chunk, Chunks, GroupBy, Group, Groups};
@@ -158,6 +159,7 @@ mod concat_impl;
 mod cons_tuples_impl;
 #[cfg(feature = "use_std")]
 mod combinations;
+mod exactly_one_err;
 mod diff;
 mod format;
 #[cfg(feature = "use_std")]
@@ -1943,13 +1945,13 @@ pub trait Itertools : Iterator {
 
     /// Return a `HashMap` of keys mapped to `Vec`s of values. Keys and values
     /// are taken from `(Key, Value)` tuple pairs yielded by the input iterator.
-    /// 
+    ///
     /// ```
     /// use itertools::Itertools;
-    /// 
+    ///
     /// let data = vec![(0, 10), (2, 12), (3, 13), (0, 20), (3, 33), (2, 42)];
     /// let lookup = data.into_iter().into_group_map();
-    /// 
+    ///
     /// assert_eq!(lookup[&0], vec![10, 20]);
     /// assert_eq!(lookup.get(&1), None);
     /// assert_eq!(lookup[&2], vec![12, 42]);
@@ -2038,6 +2040,42 @@ pub trait Itertools : Iterator {
             |x, y, _, _| Ordering::Less == compare(x, y)
         )
     }
+
+    /// If the iterator yields exactly one element, that element will be returned, otherwise
+    /// an error will be returned containing an iterator that has the same output as the input
+    /// iterator.
+    ///
+    /// This provides an additional layer of validation over just calling `Iterator::next()`.
+    /// If your assumption that there should only be one element yielded is false this provides
+    /// the opportunity to detect and handle that, preventing errors at a distance.
+    ///
+    /// # Examples
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// assert_eq!((0..10).filter(|&x| x == 2).exactly_one().unwrap(), 2);
+    /// assert!((0..10).filter(|&x| x > 1 && x < 4).exactly_one().unwrap_err().eq(2..4));
+    /// assert!((0..10).filter(|&x| x > 1 && x < 5).exactly_one().unwrap_err().eq(2..5));
+    /// assert!((0..10).filter(|&x| false).exactly_one().unwrap_err().eq(0..0));
+    /// ```
+    fn exactly_one(mut self) -> Result<Self::Item, ExactlyOneError<Self>>
+    where
+        Self: Sized,
+    {
+        match self.next() {
+            Some(first) => {
+                match self.next() {
+                    Some(second) => {
+                        Err(ExactlyOneError::new((Some(first), Some(second)), self))
+                    }
+                    None => {
+                        Ok(first)
+                    }
+                }
+            }
+            None => Err(ExactlyOneError::new((None, None), self)),
+        }
+    }
 }
 
 impl<T: ?Sized> Itertools for T where T: Iterator { }