[READ-ONLY] Mirror of https://github.com/probablykasper/csv-pipeline.
Add Flush
+250
-151
-70
src/chain.rs
-70
src/chain.rs
···
1
-
use crate::RowResult;
2
-
pub type BoxedIterator = Box<dyn Iterator<Item = RowResult>>;
3
-
4
-
/// A struct that wraps a RowResult iterator for convenience
5
-
pub struct Chain<C: Iterator<Item = RowResult>> {
6
-
iterator: C,
7
-
}
8
-
impl<C: Iterator<Item = RowResult>> Chain<C> {
9
-
pub fn new(iterator: C) -> Self {
10
-
Self { iterator: iterator }
11
-
}
12
-
pub fn with_state<S>(self, state: S) -> StatefulChainBuilder<C, S> {
13
-
StatefulChainBuilder::new(self.iterator, state)
14
-
}
15
-
}
16
-
impl<C: Iterator<Item = RowResult>> Iterator for Chain<C> {
17
-
type Item = RowResult;
18
-
19
-
fn next(&mut self) -> Option<Self::Item> {
20
-
self.iterator.next()
21
-
}
22
-
}
23
-
24
-
pub struct StatefulChainBuilder<C: Iterator<Item = RowResult>, S> {
25
-
iterator: C,
26
-
state: S,
27
-
}
28
-
impl<C: Iterator<Item = RowResult>, S> StatefulChainBuilder<C, S> {
29
-
pub fn new(iterator: C, state: S) -> Self {
30
-
Self { state, iterator }
31
-
}
32
-
pub fn map<F>(self, f: F) -> StatefulChain<C, S, F>
33
-
where
34
-
F: FnMut(RowResult, &mut S) -> RowResult,
35
-
{
36
-
StatefulChain {
37
-
iterator: self.iterator,
38
-
state: self.state,
39
-
f,
40
-
}
41
-
}
42
-
}
43
-
44
-
pub struct StatefulChain<C: Iterator<Item = RowResult>, S, F: FnMut(RowResult, &mut S) -> RowResult>
45
-
{
46
-
iterator: C,
47
-
state: S,
48
-
f: F,
49
-
}
50
-
impl<C: Iterator<Item = RowResult>, S, F> StatefulChain<C, S, F>
51
-
where
52
-
F: FnMut(RowResult, &mut S) -> RowResult,
53
-
{
54
-
pub fn into_chain(self) -> Chain<StatefulChain<C, S, F>> {
55
-
Chain::new(self)
56
-
}
57
-
}
58
-
impl<C: Iterator<Item = RowResult>, S, F> Iterator for StatefulChain<C, S, F>
59
-
where
60
-
F: FnMut(RowResult, &mut S) -> RowResult,
61
-
{
62
-
type Item = RowResult;
63
-
64
-
fn next(&mut self) -> Option<Self::Item> {
65
-
match self.iterator.next() {
66
-
Some(item) => Some((self.f)(item, &mut self.state)),
67
-
None => None,
68
-
}
69
-
}
70
-
}
+36
-23
src/lib.rs
+36
-23
src/lib.rs
···
1
-
mod chain;
2
1
mod headers;
3
2
mod pipeline;
4
-
5
-
use std::io;
3
+
mod pipeline_iterators;
4
+
pub mod target;
6
5
7
-
use csv::StringRecordsIntoIter;
8
6
pub use headers::Headers;
9
7
pub use pipeline::{Pipeline, PipelineBuilder};
10
-
11
-
#[derive(Debug, Clone, PartialEq)]
12
-
pub enum Error {
13
-
DuplicatedColumn(String),
14
-
Csv,
15
-
}
16
-
17
8
pub type Row = csv::StringRecord;
18
9
pub type RowResult = Result<Row, Error>;
19
10
20
-
pub struct RowIter<R: io::Read> {
21
-
inner: StringRecordsIntoIter<R>,
11
+
#[derive(Debug)]
12
+
pub enum Error {
13
+
Example,
14
+
Csv(csv::Error),
15
+
Io(std::io::Error),
22
16
}
23
-
impl<R: io::Read> RowIter<R> {
24
-
pub fn from_records(records: StringRecordsIntoIter<R>) -> Self {
25
-
RowIter { inner: records }
17
+
impl From<csv::Error> for Error {
18
+
fn from(error: csv::Error) -> Error {
19
+
Error::Csv(error)
20
+
}
21
+
}
22
+
23
+
impl From<std::io::Error> for Error {
24
+
fn from(error: std::io::Error) -> Error {
25
+
Error::Io(error)
26
26
}
27
27
}
28
-
impl<R: io::Read> Iterator for RowIter<R> {
29
-
type Item = RowResult;
30
28
31
-
fn next(&mut self) -> Option<Self::Item> {
32
-
self.inner.next().map(|result| {
33
-
result.map_err(|err| {
34
-
return Error::Csv;
35
-
})
29
+
#[test]
30
+
fn test_pipeline() {
31
+
let mut csv_str = String::new();
32
+
let mut pipeline = PipelineBuilder::from_path("test/Countries.csv")
33
+
.add_col("Language", |_headers, row| match row.get(1) {
34
+
Some("Norway") => Ok("Norwegian".to_string()),
35
+
_ => Ok("Unknown".to_string()),
36
36
})
37
+
.flush(target::StringTarget::new(&mut csv_str))
38
+
.build();
39
+
40
+
while let Some(err) = pipeline.next_error() {
41
+
eprintln!("Error: {err:?}");
37
42
}
43
+
drop(pipeline);
44
+
45
+
assert_eq!(
46
+
csv_str,
47
+
"ID,Country,Language\n\
48
+
1,Norway,Norwegian\n\
49
+
2,Tuvalu,Unknown\n"
50
+
);
38
51
}
+39
-58
src/pipeline.rs
+39
-58
src/pipeline.rs
···
1
1
use super::headers::Headers;
2
-
use crate::{Error, Row, RowIter, RowResult};
3
-
use csv::{Reader, ReaderBuilder};
2
+
use crate::pipeline_iterators::{AddCol, Flush};
3
+
use crate::target::Target;
4
+
use crate::{Error, Row, RowResult};
5
+
use csv::{Reader, ReaderBuilder, StringRecordsIntoIter};
4
6
use std::fs::File;
7
+
use std::io;
5
8
use std::path::Path;
6
9
7
10
pub struct PipelineBuilder<'a> {
···
51
54
F: FnMut(&Headers, &Row) -> Result<String, Error> + 'a,
52
55
{
53
56
self.headers.push_field(name);
54
-
55
-
let add_col = AddCol {
57
+
self.iterator = Box::new(AddCol {
56
58
iterator: self.iterator,
57
59
f: get_value,
58
60
headers: self.headers.clone(),
59
-
};
61
+
});
62
+
self
63
+
}
60
64
61
-
self.iterator = Box::new(add_col);
62
-
65
+
pub fn flush(mut self, target: impl Target + 'a) -> Self {
66
+
let flush = Flush::new(self.iterator, target, self.headers.clone());
67
+
self.iterator = Box::new(flush);
63
68
self
64
69
}
65
70
···
71
76
}
72
77
}
73
78
74
-
struct AddCol<I, F: FnMut(&Headers, &Row) -> Result<String, Error>> {
75
-
iterator: I,
76
-
f: F,
77
-
headers: Headers,
79
+
pub struct Pipeline<'a> {
80
+
pub headers: Headers,
81
+
pub iterator: Box<dyn Iterator<Item = RowResult> + 'a>,
78
82
}
79
-
impl<I, F> Iterator for AddCol<I, F>
80
-
where
81
-
I: Iterator<Item = RowResult>,
82
-
F: FnMut(&Headers, &Row) -> Result<String, Error>,
83
-
{
84
-
type Item = RowResult;
85
83
86
-
fn next(&mut self) -> Option<Self::Item> {
87
-
let row = match self.iterator.next() {
88
-
Some(Ok(row)) => row,
89
-
Some(Err(e)) => return Some(Err(e)),
90
-
None => return None,
91
-
};
92
-
match (self.f)(&self.headers, &row) {
93
-
Ok(value) => {
94
-
let mut row = row;
95
-
row.push_field(&value);
96
-
Some(Ok(row))
84
+
impl<'a> Pipeline<'a> {
85
+
/// Advances the iterator until an error is found.
86
+
///
87
+
/// Returns `None` when the iterator is finished.
88
+
pub fn next_error(&mut self) -> Option<Error> {
89
+
while let Some(item) = self.next() {
90
+
if let Err(err) = item {
91
+
return Some(err);
97
92
}
98
-
Err(e) => Some(Err(e)),
99
93
}
94
+
None
100
95
}
101
96
}
102
-
103
-
pub struct Pipeline<'a> {
104
-
pub headers: Headers,
105
-
pub iterator: Box<dyn Iterator<Item = RowResult> + 'a>,
106
-
}
107
97
impl<'a> Iterator for Pipeline<'a> {
108
98
type Item = RowResult;
109
99
···
112
102
}
113
103
}
114
104
115
-
#[cfg(test)]
116
-
mod tests {
117
-
use crate::PipelineBuilder;
105
+
pub struct RowIter<R: io::Read> {
106
+
inner: StringRecordsIntoIter<R>,
107
+
}
108
+
impl<R: io::Read> RowIter<R> {
109
+
pub fn from_records(records: StringRecordsIntoIter<R>) -> Self {
110
+
RowIter { inner: records }
111
+
}
112
+
}
113
+
impl<R: io::Read> Iterator for RowIter<R> {
114
+
type Item = RowResult;
118
115
119
-
#[test]
120
-
fn add_col() {
121
-
let mut pipeline = PipelineBuilder::from_path("test/Countries.csv")
122
-
.add_col("Language", |_headers, row| match row.get(1) {
123
-
Some("Norway") => Ok("Norwegian".to_string()),
124
-
_ => Ok("Unknown".to_string()),
116
+
fn next(&mut self) -> Option<Self::Item> {
117
+
self.inner.next().map(|result| {
118
+
result.map_err(|err| {
119
+
return Error::from(err);
125
120
})
126
-
.build();
127
-
128
-
let mut writer = csv::Writer::from_writer(vec![]);
129
-
writer.write_record(&pipeline.headers).unwrap();
130
-
while let Some(item) = pipeline.next() {
131
-
writer.write_record(&item.unwrap()).unwrap();
132
-
}
133
-
let csv_str = String::from_utf8(writer.into_inner().unwrap()).unwrap();
134
-
135
-
assert_eq!(
136
-
csv_str,
137
-
"ID,Country,Language\n\
138
-
1,Norway,Norwegian\n\
139
-
2,Tuvalu,Unknown\n"
140
-
);
121
+
})
141
122
}
142
123
}
+74
src/pipeline_iterators.rs
+74
src/pipeline_iterators.rs
···
1
+
use super::headers::Headers;
2
+
use crate::target::Target;
3
+
use crate::{Error, Row, RowResult};
4
+
5
+
pub struct AddCol<I, F: FnMut(&Headers, &Row) -> Result<String, Error>> {
6
+
pub iterator: I,
7
+
pub f: F,
8
+
pub headers: Headers,
9
+
}
10
+
impl<I, F> Iterator for AddCol<I, F>
11
+
where
12
+
I: Iterator<Item = RowResult>,
13
+
F: FnMut(&Headers, &Row) -> Result<String, Error>,
14
+
{
15
+
type Item = RowResult;
16
+
17
+
fn next(&mut self) -> Option<Self::Item> {
18
+
let row = match self.iterator.next()? {
19
+
Ok(row) => row,
20
+
Err(e) => return Some(Err(e)),
21
+
};
22
+
match (self.f)(&self.headers, &row) {
23
+
Ok(value) => {
24
+
let mut row = row;
25
+
row.push_field(&value);
26
+
Some(Ok(row))
27
+
}
28
+
Err(e) => Some(Err(e)),
29
+
}
30
+
}
31
+
}
32
+
33
+
pub struct Flush<I, T> {
34
+
pub iterator: I,
35
+
pub target: T,
36
+
/// `None` if headers have been written, `Some` otherwise
37
+
headers: Option<Headers>,
38
+
}
39
+
40
+
impl<I, T> Flush<I, T> {
41
+
pub fn new(iterator: I, target: T, headers: Headers) -> Self {
42
+
Self {
43
+
iterator,
44
+
target,
45
+
headers: Some(headers),
46
+
}
47
+
}
48
+
}
49
+
impl<I, T> Iterator for Flush<I, T>
50
+
where
51
+
I: Iterator<Item = RowResult>,
52
+
T: Target,
53
+
{
54
+
type Item = RowResult;
55
+
56
+
fn next(&mut self) -> Option<Self::Item> {
57
+
if let Some(headers) = &self.headers {
58
+
match self.target.write_headers(headers) {
59
+
Ok(()) => self.headers = None,
60
+
Err(e) => return Some(Err(e)),
61
+
}
62
+
}
63
+
64
+
let row = match self.iterator.next()? {
65
+
Ok(row) => row,
66
+
Err(e) => return Some(Err(e)),
67
+
};
68
+
let r = match self.target.write_row(&row) {
69
+
Ok(()) => Some(Ok(row)),
70
+
Err(e) => Some(Err(e)),
71
+
};
72
+
r
73
+
}
74
+
}
+101
src/target.rs
+101
src/target.rs
···
1
+
use crate::{Error, Headers, Row};
2
+
use csv::WriterBuilder;
3
+
use std::fs::{self, File};
4
+
use std::io;
5
+
use std::path::PathBuf;
6
+
7
+
pub trait Target {
8
+
/// Useful for initializations
9
+
fn write_headers(&mut self, headers: &Headers) -> Result<(), Error>;
10
+
fn write_row(&mut self, row: &Row) -> Result<(), Error>;
11
+
}
12
+
13
+
pub struct PathTarget {
14
+
path: PathBuf,
15
+
writer: Option<csv::Writer<File>>,
16
+
}
17
+
impl PathTarget {
18
+
pub fn new<P: Into<PathBuf>>(path: P) -> Self {
19
+
Self {
20
+
path: path.into(),
21
+
writer: None,
22
+
}
23
+
}
24
+
}
25
+
impl Target for PathTarget {
26
+
fn write_headers(&mut self, headers: &Headers) -> Result<(), Error> {
27
+
if let Some(parent) = self.path.parent() {
28
+
fs::create_dir_all(parent)?;
29
+
}
30
+
31
+
self.writer = Some(csv::Writer::from_path(&self.path)?);
32
+
self.write_row(headers.get_row())
33
+
}
34
+
fn write_row(&mut self, row: &Row) -> Result<(), Error> {
35
+
self.writer.as_mut().unwrap().write_record(row)?;
36
+
Ok(())
37
+
}
38
+
}
39
+
40
+
pub struct StdoutTarget {
41
+
writer: Option<csv::Writer<io::Stdout>>,
42
+
}
43
+
impl StdoutTarget {
44
+
pub fn new() -> Self {
45
+
Self { writer: None }
46
+
}
47
+
}
48
+
impl Target for StdoutTarget {
49
+
fn write_headers(&mut self, headers: &Headers) -> Result<(), Error> {
50
+
let writer = WriterBuilder::new().from_writer(io::stdout());
51
+
self.writer = Some(writer);
52
+
self.write_row(headers.get_row())?;
53
+
Ok(())
54
+
}
55
+
fn write_row(&mut self, row: &Row) -> Result<(), Error> {
56
+
self.writer.as_mut().unwrap().write_record(row)?;
57
+
Ok(())
58
+
}
59
+
}
60
+
61
+
pub struct StderrTarget {
62
+
writer: Option<csv::Writer<io::Stderr>>,
63
+
}
64
+
impl StderrTarget {
65
+
pub fn new() -> Self {
66
+
Self { writer: None }
67
+
}
68
+
}
69
+
impl Target for StderrTarget {
70
+
fn write_headers(&mut self, headers: &Headers) -> Result<(), Error> {
71
+
let writer = WriterBuilder::new().from_writer(io::stderr());
72
+
self.writer = Some(writer);
73
+
self.write_row(headers.get_row())
74
+
}
75
+
fn write_row(&mut self, row: &Row) -> Result<(), Error> {
76
+
self.writer.as_mut().unwrap().write_record(row)?;
77
+
Ok(())
78
+
}
79
+
}
80
+
81
+
pub struct StringTarget<'a> {
82
+
s: &'a mut String,
83
+
}
84
+
impl<'a> StringTarget<'a> {
85
+
pub fn new(s: &'a mut String) -> Self {
86
+
Self { s }
87
+
}
88
+
}
89
+
impl<'a> Target for StringTarget<'a> {
90
+
fn write_headers(&mut self, headers: &Headers) -> Result<(), Error> {
91
+
self.write_row(headers.get_row())
92
+
}
93
+
fn write_row(&mut self, row: &Row) -> Result<(), Error> {
94
+
let mut wtr = csv::Writer::from_writer(vec![]);
95
+
wtr.write_record(row)?;
96
+
// not sure if this unwrap is safe
97
+
let row_s = String::from_utf8(wtr.into_inner().unwrap()).unwrap();
98
+
*self.s += &row_s;
99
+
Ok(())
100
+
}
101
+
}