use crate::{
Headers, Row, RowStream,
error::{Error, RowResult},
};
pub struct Select<'a, I> {
iter: I,
columns: Vec<&'a str>,
headers: Headers,
old_headers: Headers,
}
impl<'a, I> Select<'a, I>
where
I: RowStream,
{
pub fn new(iter: I, columns: Vec<&'a str>) -> Select<I> {
let headers: Headers = Row::from(columns.clone()).into();
Select{
old_headers: iter.headers().clone(),
iter,
columns,
headers,
}
}
}
pub struct IntoIter<'a, I> {
iter: I,
columns: Vec<&'a str>,
old_headers: Headers,
}
impl<'a, I> Iterator for IntoIter<'a, I>
where
I: Iterator<Item = RowResult>,
{
type Item = RowResult;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next().map(|result| {
result.and_then(|val| {
let values = self.columns.iter()
.map(|header| {
self.old_headers
.get_field(&val, header)
.ok_or_else(|| Error::ColumnNotFound(header.to_string()))
})
.collect::<Result<Vec<_>, Error>>()?;
Ok(values.into())
})
})
}
}
impl<'a, I> IntoIterator for Select<'a, I>
where
I: RowStream,
{
type Item = RowResult;
type IntoIter = IntoIter<'a, I::IntoIter>;
fn into_iter(self) -> Self::IntoIter {
Self::IntoIter {
iter: self.iter.into_iter(),
columns: self.columns,
old_headers: self.old_headers,
}
}
}
impl<'a, I> RowStream for Select<'a, I>
where
I: RowStream,
{
fn headers(&self) -> &Headers {
&self.headers
}
}
#[cfg(test)]
mod tests {
use super::Select;
use crate::mock::MockStream;
use crate::{Row, RowStream};
#[test]
fn keeps_the_correct_headers() {
let iter = MockStream::from_rows(
vec![
Ok(Row::from(vec!["id", "val", "path"])),
Ok(Row::from(vec!["1", "40", "/tmp/a1m.csv"])),
Ok(Row::from(vec!["2", "39", "/tmp/a1m.csv"])),
Ok(Row::from(vec!["3", "38", "/tmp/a2m.csv"])),
Ok(Row::from(vec!["4", "37", "/tmp/a2m.csv"])),
]
.into_iter(),
)
.unwrap();
let del = Select::new(
iter,
vec!["id", "val"],
);
assert_eq!(
*del.headers(),
Row::from(vec!["id", "val"]).into(),
);
let mut del = del.into_iter();
assert_eq!(
del.next().unwrap().unwrap(),
Row::from(vec!["1", "40"])
);
assert_eq!(
del.next().unwrap().unwrap(),
Row::from(vec!["2", "39"])
);
assert_eq!(
del.next().unwrap().unwrap(),
Row::from(vec!["3", "38"])
);
assert_eq!(
del.next().unwrap().unwrap(),
Row::from(vec!["4", "37"])
);
}
#[test]
fn keeps_them_in_the_correct_order() {
let iter = MockStream::from_rows(
vec![
Ok(Row::from(vec!["id", "val", "path"])),
Ok(Row::from(vec!["1", "40", "/tmp/a1m.csv"])),
Ok(Row::from(vec!["2", "39", "/tmp/a1m.csv"])),
Ok(Row::from(vec!["3", "38", "/tmp/a2m.csv"])),
Ok(Row::from(vec!["4", "37", "/tmp/a2m.csv"])),
]
.into_iter(),
)
.unwrap();
let sel = Select::new(
iter,
vec!["path", "val"],
);
assert_eq!(
*sel.headers(),
Row::from(vec!["path", "val"]).into(),
);
let mut sel = sel.into_iter();
assert_eq!(
sel.next().unwrap().unwrap(),
Row::from(vec!["/tmp/a1m.csv", "40"])
);
assert_eq!(
sel.next().unwrap().unwrap(),
Row::from(vec!["/tmp/a1m.csv", "39"])
);
assert_eq!(
sel.next().unwrap().unwrap(),
Row::from(vec!["/tmp/a2m.csv", "38"])
);
assert_eq!(
sel.next().unwrap().unwrap(),
Row::from(vec!["/tmp/a2m.csv", "37"])
);
}
}