diff --git a/src/exactly_one_err.rs b/src/exactly_one_err.rs new file mode 100644 index 000000000..51cfbc347 --- /dev/null +++ b/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 +where + I: Iterator, +{ + first_two: (Option, Option), + inner: I, +} + +impl ExactlyOneError +where + I: Iterator, +{ + /// Creates a new `ExactlyOneErr` iterator. + pub fn new(first_two: (Option, Option), inner: I) -> Self { + Self { first_two, inner } + } +} + +impl Iterator for ExactlyOneError +where + I: Iterator, +{ + type Item = I::Item; + + fn next(&mut self) -> Option { + self.first_two + .0 + .take() + .or_else(|| self.first_two.1.take()) + .or_else(|| self.inner.next()) + } + + fn size_hint(&self) -> (usize, Option) { + 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 ExactSizeIterator for ExactlyOneError where I: ExactSizeIterator {} diff --git a/src/lib.rs b/src/lib.rs index c6518ff60..fdedf2b0a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -74,6 +74,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}; @@ -128,6 +129,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")] @@ -1888,13 +1890,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]); @@ -1983,6 +1985,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> + 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 Itertools for T where T: Iterator { }