1use crate::{ Headers, RowStream, error };
3
4mod colspec;
5mod column;
6
7pub use colspec::ColSpec;
8pub use column::Column;
9
10#[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}