csvsc/
add.rs

1//!Machinery for adding columns.
2use crate::{ Headers, RowStream, error };
3
4mod colspec;
5mod column;
6
7pub use colspec::ColSpec;
8pub use column::Column;
9
10/// Adds a column to each register. It can be based on existing ones
11/// or the source filename.
12#[derive(Debug)]
13pub struct Add<I> {
14    iter: I,
15    column: ColSpec,
16    headers: Headers,
17}
18
19impl<I> Add<I>
20where
21    I: RowStream,
22{
23    pub fn new(iter: I, column: ColSpec) -> error::Result<Add<I>> {
24        let mut headers = iter.headers().clone();
25        let colname = match &column {
26            ColSpec::Regex { colname, .. } => colname,
27            ColSpec::Mix { colname, .. } => colname,
28        };
29
30        if !headers.add(colname) {
31            return Err(error::Error::DuplicatedColumn(colname.to_string()));
32        }
33
34        Ok(Add{
35            iter,
36            column,
37            headers,
38        })
39    }
40}
41
42pub struct IntoIter<I> {
43    iter: I,
44    column: ColSpec,
45    headers: Headers,
46}
47
48impl<I> Iterator for IntoIter<I>
49where
50    I: Iterator<Item = error::RowResult>,
51{
52    type Item = error::RowResult;
53
54    fn next(&mut self) -> Option<Self::Item> {
55        self.iter.next().map(|result| {
56            result.and_then(|mut val| {
57                match self.column.compute(&val, &self.headers) {
58                    Ok(s) => val.push_field(&s),
59                    Err(e) => return Err(e),
60                }
61
62                Ok(val)
63            })
64        })
65    }
66}
67
68impl<I> IntoIterator for Add<I>
69where
70    I: RowStream,
71{
72    type Item = error::RowResult;
73
74    type IntoIter = IntoIter<I::IntoIter>;
75
76    fn into_iter(self) -> Self::IntoIter {
77        Self::IntoIter {
78            iter: self.iter.into_iter(),
79            column: self.column,
80            headers: self.headers,
81        }
82    }
83}
84
85impl<I> RowStream for Add<I>
86where
87    I: RowStream,
88{
89    fn headers(&self) -> &Headers {
90        &self.headers
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::{Add, RowStream, ColSpec};
97    use crate::{
98        Row,
99        mock::MockStream,
100        error::Error,
101    };
102    use regex::Regex;
103
104    #[test]
105    fn test_add() {
106        let iter = MockStream::from_rows(
107            vec![
108                Ok(Row::from(vec!["id", "val", "path"])),
109                Ok(Row::from(vec!["1", "40", "/tmp/a1m.csv"])),
110                Ok(Row::from(vec!["2", "39", "/tmp/a1m.csv"])),
111                Ok(Row::from(vec!["3", "38", "/tmp/a2m.csv"])),
112                Ok(Row::from(vec!["4", "37", "/tmp/a2m.csv"])),
113            ]
114            .into_iter(),
115        )
116        .unwrap();
117
118        let add = Add::new(
119            iter,
120            ColSpec::Regex {
121                source: "path".to_string(),
122                colname: "new".to_string(),
123                coldef: "$1".to_string(),
124                regex: Regex::new("a([0-9]+)m\\.csv$").unwrap(),
125            },
126        ).unwrap();
127
128        assert_eq!(
129            *add.headers(),
130            Row::from(vec!["id", "val", "path", "new"]).into(),
131        );
132
133        let mut add = add.into_iter();
134
135        assert_eq!(
136            add.next().unwrap().unwrap(),
137            Row::from(vec!["1", "40", "/tmp/a1m.csv", "1"])
138        );
139        assert_eq!(
140            add.next().unwrap().unwrap(),
141            Row::from(vec!["2", "39", "/tmp/a1m.csv", "1"])
142        );
143        assert_eq!(
144            add.next().unwrap().unwrap(),
145            Row::from(vec!["3", "38", "/tmp/a2m.csv", "2"])
146        );
147        assert_eq!(
148            add.next().unwrap().unwrap(),
149            Row::from(vec!["4", "37", "/tmp/a2m.csv", "2"])
150        );
151    }
152
153    #[test]
154    fn test_add_doesnt_swallow_errors() {
155        let iter = MockStream::from_rows(
156            vec![
157                Ok(Row::from(vec!["a"])),
158                Ok(Row::from(vec!["1"])),
159                Err(Error::InconsistentHeaders),
160                Ok(Row::from(vec!["3"])),
161            ]
162            .into_iter(),
163        )
164        .unwrap();
165
166        let add = Add::new(
167            iter,
168            ColSpec::Mix {
169                colname: "b".to_string(),
170                coldef: "1".to_string(),
171            },
172        ).unwrap();
173
174        assert_eq!(
175            *add.headers(),
176            Row::from(vec!["a", "b"]).into(),
177        );
178
179        let mut add = add.into_iter();
180
181        assert_eq!(
182            add.next().unwrap().unwrap(),
183            Row::from(vec!["1", "1"])
184        );
185
186        match add.next() {
187            Some(Err(Error::InconsistentHeaders)) => {},
188            _ => unreachable!(),
189        }
190
191        assert_eq!(
192            add.next().unwrap().unwrap(),
193            Row::from(vec!["3", "1"])
194        );
195    }
196}