more_itertools/grouping/
chunked.rs

1use std::mem::swap;
2
3use crate::error::Error;
4use crate::error;
5
6pub struct Chunked<T> {
7    iter: Box<dyn Iterator<Item = T>>,
8    buf: Vec<T>,
9    n: usize,
10    strict: bool
11}
12
13
14impl<T> Iterator for Chunked<T> {
15    type Item = Result<Vec<T>, Error>;
16
17    fn next(&mut self) -> Option<Self::Item> {
18        if self.n == 0 {
19            return Some(Err(error::value_error(String::from("n should not be 0"))));
20        }
21
22        let mut ret = Vec::new();
23        while self.buf.len() < self.n {
24            let ele = self.iter.next();
25            match ele {
26                Some(ele_value) => { self.buf.push(ele_value) }
27                None => {
28                    if self.buf.len() == 0 {
29                        return None;
30                    } else {
31                        swap(&mut ret, &mut self.buf);
32                        self.buf = Vec::new();
33                        if self.strict && ret.len() < self.n {
34                            return Some(Err(error::value_error(String::from("iterable is not divisible by n."))));
35                        }
36                        return Some(Ok(ret));
37                    }
38                }
39            }
40
41            if self.buf.len() == self.n {
42                ret = Vec::new();
43                swap(&mut ret, &mut self.buf);
44                return Some(Ok(ret));
45            }
46        }
47
48        return Some(Ok(ret));
49
50    }
51}
52
53/// https://more-itertools.readthedocs.io/en/v10.2.0/_modules/more_itertools/more.html#chunked
54pub fn chunked<T>(iter: Box<dyn Iterator<Item = T>>, n: usize, strict: bool) -> Box<dyn Iterator<Item = Result<Vec<T>, Error>>>
55where
56    T: 'static,
57{
58    Box::new(Chunked {
59        iter,
60        buf: Vec::new(),
61        n,
62        strict
63    })
64}
65
66
67#[cfg(test)]
68mod tests {
69    use crate::itertools::iter::iter_from_vec;
70
71    use super::*;
72
73    #[test]
74    fn test1_no_strict() {
75        let mut it = chunked(iter_from_vec(vec![1,2,3,4,5,6,7,8,9,10]), 3, false);
76
77        assert_eq!(vec![1,2,3], it.next().unwrap().ok().unwrap());
78        assert_eq!(vec![4,5,6], it.next().unwrap().ok().unwrap());
79        assert_eq!(vec![7,8,9], it.next().unwrap().ok().unwrap());
80        assert_eq!(vec![10], it.next().unwrap().ok().unwrap());
81        assert_eq!(None, it.next());
82        assert_eq!(None, it.next());
83    }
84
85    #[test]
86    fn test2_strict() {
87        let mut it = chunked(iter_from_vec(vec![1,2,3,4,5,6,7,8,9,10]), 3, true);
88
89        assert_eq!(vec![1,2,3], it.next().unwrap().ok().unwrap());
90        assert_eq!(vec![4,5,6], it.next().unwrap().ok().unwrap());
91        assert_eq!(vec![7,8,9], it.next().unwrap().ok().unwrap());
92
93        assert_eq!(error::Kind::ValueError, it.next().unwrap().err().unwrap().kind());
94    }
95
96    #[test]
97    fn test3_no_strict_chars() {
98        let mut it = chunked(iter_from_vec("abcdefghij".chars().collect()), 3, false);
99
100        assert_eq!(vec!['a', 'b', 'c'], it.next().unwrap().ok().unwrap());
101        assert_eq!(vec!['d', 'e', 'f'], it.next().unwrap().ok().unwrap());
102        assert_eq!(vec!['g', 'h', 'i'], it.next().unwrap().ok().unwrap());
103        assert_eq!(vec!['j'], it.next().unwrap().ok().unwrap());
104
105        assert_eq!(None, it.next());
106        assert_eq!(None, it.next());
107    }
108
109    #[test]
110    fn test3_no_strict_string() {
111        let v = vec![String::from("1"),
112                            String::from("2"),
113                            String::from("3"),
114                            String::from("4")
115                            ];
116
117        let mut it = chunked(iter_from_vec(v), 3, false);
118
119        assert_eq!(vec![String::from("1"), String::from("2") ,String::from("3")], it.next().unwrap().ok().unwrap());
120        assert_eq!(vec![String::from("4")], it.next().unwrap().ok().unwrap());
121
122        assert_eq!(None, it.next());
123        assert_eq!(None, it.next());
124    }
125
126    #[test]
127    fn test4_value_error_n_is_0() {
128        let mut it = chunked(iter_from_vec(vec![1,2,3,4,5,6,7,8,9,10]), 0, false);
129        assert_eq!(error::Kind::ValueError, it.next().unwrap().err().unwrap().kind());
130    }
131}
132
more_itertools/grouping/chunked.rs

more_itertools/grouping/
chunked.rs
