VecOrStr implementing IntoIterator<Item = String>

vecorstr.rs

use std::vec::IntoIter;

pub enum VecOrStr {
    Vec { inner: Vec<String> },
    Str { inner: String },
}

impl VecOrStr {
    pub fn from_iter<S: AsRef<str>>(iter: impl IntoIterator<Item = S>) -> Self {
        let inner = iter.into_iter().map(|x| x.as_ref().to_owned()).collect();
        Self::Vec { inner }
    }

    pub fn from_str<S: AsRef<str>>(s: S) -> Self {
        let inner = s.as_ref().to_owned();
        Self::Str { inner }
    }
}

pub enum VecOrStrIter {
    Vec { inner: IntoIter<String> },
    Str { inner: Option<String> },
}

impl Iterator for VecOrStrIter {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Vec { inner } => inner.next(),
            Self::Str { inner } => inner.take(),
        }
    }
}

impl IntoIterator for VecOrStr {
    type Item = String;
    type IntoIter = VecOrStrIter;

    fn into_iter(self) -> Self::IntoIter {
        match self {
            Self::Vec { inner } => VecOrStrIter::Vec {
                inner: inner.into_iter(),
            },
            Self::Str { inner } => VecOrStrIter::Str { inner: Some(inner) },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_works() {
        let vs1 = VecOrStr::from_iter(vec!["1", "2"]);
        assert_eq!(vec!["1", "2"], vs1.into_iter().collect::<Vec<_>>());

        let vs2 = VecOrStr::from_str("3");
        assert_eq!(vec!["3"], vs2.into_iter().collect::<Vec<_>>());
    }
}