csv_pipeline/
target.rs

1use crate::{Headers, Row};
2use csv::WriterBuilder;
3use std::fs::{self, File};
4use std::io;
5use std::path::PathBuf;
6
7pub trait Target {
8	/// Useful for initializations
9	fn write_headers(&mut self, headers: &Headers) -> Result<(), csv::Error>;
10	fn write_row(&mut self, row: &Row) -> Result<(), csv::Error>;
11}
12
13pub struct PathTarget {
14	path: PathBuf,
15	writer: Option<csv::Writer<File>>,
16}
17impl PathTarget {
18	pub fn new<P: Into<PathBuf>>(path: P) -> Self {
19		Self {
20			path: path.into(),
21			writer: None,
22		}
23	}
24}
25impl Target for PathTarget {
26	fn write_headers(&mut self, headers: &Headers) -> Result<(), csv::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<(), csv::Error> {
35		self.writer.as_mut().unwrap().write_record(row)?;
36		Ok(())
37	}
38}
39
40pub struct StdoutTarget {
41	writer: Option<csv::Writer<io::Stdout>>,
42}
43impl StdoutTarget {
44	pub fn new() -> Self {
45		Self { writer: None }
46	}
47}
48impl Target for StdoutTarget {
49	fn write_headers(&mut self, headers: &Headers) -> Result<(), csv::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<(), csv::Error> {
56		self.writer.as_mut().unwrap().write_record(row)?;
57		Ok(())
58	}
59}
60
61pub struct StderrTarget {
62	writer: Option<csv::Writer<io::Stderr>>,
63}
64impl StderrTarget {
65	pub fn new() -> Self {
66		Self { writer: None }
67	}
68}
69impl Target for StderrTarget {
70	fn write_headers(&mut self, headers: &Headers) -> Result<(), csv::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<(), csv::Error> {
76		self.writer.as_mut().unwrap().write_record(row)?;
77		Ok(())
78	}
79}
80
81pub struct StringWriter<'a> {
82	s: &'a mut String,
83}
84impl<'a> io::Write for StringWriter<'a> {
85	fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
86		match std::str::from_utf8(buf) {
87			Ok(s) => {
88				self.s.push_str(s);
89				Ok(buf.len())
90			}
91			Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e)),
92		}
93	}
94	fn flush(&mut self) -> io::Result<()> {
95		Ok(())
96	}
97}
98
99pub struct StringTarget<'a> {
100	writer: csv::Writer<StringWriter<'a>>,
101}
102impl<'a> StringTarget<'a> {
103	pub fn new(s: &'a mut String) -> Self {
104		let writer = WriterBuilder::new().from_writer(StringWriter { s });
105		Self { writer }
106	}
107}
108impl<'a> Target for StringTarget<'a> {
109	fn write_headers(&mut self, headers: &Headers) -> Result<(), csv::Error> {
110		self.write_row(headers.get_row())
111	}
112	fn write_row(&mut self, row: &Row) -> Result<(), csv::Error> {
113		self.writer.write_record(row)?;
114		Ok(())
115	}
116}