use crate::{Headers, RowStream, RowResult, Error};
pub struct FilterCol<I, F> {
iter: I,
f: F,
colname: String,
headers: Headers,
}
impl<I, F> FilterCol<I, F>
where
I: RowStream,
F: Fn(&str) -> bool,
{
pub fn new(
iter: I,
colname: String,
f: F,
) -> Result<FilterCol<I, F>, Error> {
let headers = iter.headers().clone();
if !headers.contains_key(&colname) {
return Err(Error::ColumnNotFound(colname));
}
Ok(FilterCol {
iter,
f,
colname,
headers,
})
}
}
pub struct IntoIter<I, F> {
iter: I,
f: F,
colname: String,
headers: Headers,
}
impl<I, F> Iterator for IntoIter<I, F>
where
I: Iterator<Item = RowResult>,
F: Fn(&str) -> bool,
{
type Item = RowResult;
fn next(&mut self) -> Option<Self::Item> {
for rs in self.iter.by_ref() {
match rs {
Ok(row) => if let Some(val) = self.headers.get_field(&row, &self.colname) {
if (self.f)(val) {
return Some(Ok(row));
}
} else {
return Some(Err(Error::ColumnNotFound(self.colname.clone())));
}
Err(e) => return Some(Err(e)),
}
}
None
}
}
impl<I, F> IntoIterator for FilterCol<I, F>
where
I: RowStream,
F: Fn(&str) -> bool,
{
type Item = RowResult;
type IntoIter = IntoIter<I::IntoIter, F>;
fn into_iter(self) -> Self::IntoIter {
Self::IntoIter {
iter: self.iter.into_iter(),
f: self.f,
colname: self.colname,
headers: self.headers,
}
}
}
impl<I, F> RowStream for FilterCol<I, F>
where
I: RowStream,
F: Fn(&str) -> bool,
{
fn headers(&self) -> &Headers {
&self.headers
}
}
#[cfg(test)]
mod tests {
use super::FilterCol;
use crate::{Row, MockStream, error};
#[test]
fn test_filter() {
let iter = MockStream::from_rows(
vec![
Ok(Row::from(vec!["a", "b"])),
Ok(Row::from(vec!["1", "2"])),
Ok(Row::from(vec!["2", "7"])),
Err(error::Error::InconsistentHeaders),
Ok(Row::from(vec!["1", "1"])),
Ok(Row::from(vec!["2", "9"])),
]
.into_iter(),
)
.unwrap();
let mut r = FilterCol::new(iter, "b".into(), |value| {
value.parse::<i32>().unwrap() > 2
}).unwrap()
.into_iter();
assert!(r.next().unwrap().is_ok());
assert!(r.next().unwrap().is_err());
assert!(r.next().unwrap().is_ok());
assert!(r.next().is_none());
}
}