use std::io::Error;
use csv::ByteRecord;
use resext::{CtxResult, ResExt};
use serde::Serialize;
pub(crate) enum WriterStreams<I>
where
I: Iterator<Item = CtxResult<DataTypes, Error>>,
{
Values { iter: I },
Table { headers: Vec<String>, iter: I },
Ndjson { values: I },
}
pub(crate) enum DataTypes {
Json(serde_json::Value),
Toml(toml::Value),
Csv(ByteRecord),
}
impl Serialize for DataTypes {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
match self {
DataTypes::Json(j) => j.serialize(serializer),
DataTypes::Toml(t) => t.serialize(serializer),
DataTypes::Csv(_) => unreachable!(),
}
}
}
pub(crate) fn into_byte_record(brec: CtxResult<DataTypes, Error>) -> CtxResult<ByteRecord, Error> {
match brec.context("Failed to unwrap record")? {
DataTypes::Csv(csv) => Ok(csv),
_ => unreachable!(),
}
}
static NEEDS_ESCAPE: [bool; 256] = {
let mut table = [false; 256];
table[b'\\' as usize] = true;
table[b'"' as usize] = true;
table[b'\n' as usize] = true;
table[b'\r' as usize] = true;
table[b'\t' as usize] = true;
table
};
#[inline]
pub(crate) fn escape(byte: u8, output: &mut Vec<u8>) {
if NEEDS_ESCAPE[byte as usize] {
output.reserve_exact(2);
output.push(b'\\');
match byte {
b'\\' => output.push(b'\\'),
b'"' => output.push(b'"'),
b'\n' => output.push(b'n'),
b'\r' => output.push(b'r'),
b'\t' => output.push(b't'),
_ => unreachable!(),
}
} else {
output.push(byte);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_escape_quote() {
let mut out = Vec::new();
escape(b'"', &mut out);
assert_eq!(out, b"\\\"");
}
#[test]
fn test_escape_newline() {
let mut out = Vec::new();
escape(b'\n', &mut out);
assert_eq!(out, b"\\n");
}
#[test]
fn test_escape_backslash() {
let mut out = Vec::new();
escape(b'\\', &mut out);
assert_eq!(out, b"\\\\");
}
#[test]
fn test_escape_tab() {
let mut out = Vec::new();
escape(b'\t', &mut out);
assert_eq!(out, b"\\t");
}
#[test]
fn test_escape_carriage_return() {
let mut out = Vec::new();
escape(b'\r', &mut out);
assert_eq!(out, b"\\r");
}
#[test]
fn test_no_escape_needed() {
let mut out = Vec::new();
escape(b'a', &mut out);
assert_eq!(out, b"a");
out.clear();
escape(b'0', &mut out);
assert_eq!(out, b"0");
}
#[test]
fn test_needs_escape_table() {
assert!(NEEDS_ESCAPE[b'"' as usize]);
assert!(NEEDS_ESCAPE[b'\n' as usize]);
assert!(NEEDS_ESCAPE[b'\r' as usize]);
assert!(NEEDS_ESCAPE[b'\t' as usize]);
assert!(NEEDS_ESCAPE[b'\\' as usize]);
assert!(!NEEDS_ESCAPE[b'a' as usize]);
assert!(!NEEDS_ESCAPE[b'0' as usize]);
assert!(!NEEDS_ESCAPE[b' ' as usize]);
assert!(!NEEDS_ESCAPE[b',' as usize]);
}
}