use crate::prelude::*;
use crate::*;
use memchr::{memchr2, memchr3};
use read_line::BackslashMode;
use read_line::Quotes;
use read_line::ReadResult;
use std::collections::HashSet;
use std::io;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Escape {
QuoteDoubled(u8, u8), AlwaysQuoteDoubled(u8, u8), QuoteBackslash(u8, u8), AlwaysQuoteBackslash(u8, u8), Backslash,
Error,
Trust,
Replace(u8),
Delete,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Header {
Cdx,
Yes,
No,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Config {
pub delimiter: u8,
pub escape: Escape,
pub header: Header,
}
impl Default for Config {
fn default() -> Self {
Self { delimiter: b'\t', escape: Escape::Replace(b' '), header: Header::Yes }
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Eol {
Plain(ReadResult),
Fancy(Vec<u8>),
}
impl Default for Eol {
fn default() -> Self {
Self::Plain(ReadResult::Lf)
}
}
impl Eol {
fn as_bytes(&self) -> &[u8] {
match self {
Self::Plain(r) => r.as_bytes(),
Self::Fancy(v) => v,
}
}
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Spec {
pub delimiter: Option<u8>,
pub escape: Option<Escape>,
pub header: Option<Header>,
pub eol: Option<Eol>,
}
#[derive(Debug, Default)]
pub struct LineWriter {
config: Config,
eol: Vec<u8>,
record_buf: Vec<u8>,
column_buf: Vec<u8>,
column_open: bool,
wrote_any_column: bool,
}
#[inline]
fn has_delimiter_or_newline(data: &[u8], delimiter: u8) -> bool {
find_replace_target(data, delimiter).is_some()
}
#[inline]
fn find_close_quote_target(haystack: &[u8], _open: u8, close: u8) -> Option<usize> {
memchr::memchr(close, haystack)
}
fn write_quoted_doubled(data: &[u8], out: &mut Vec<u8>, open: u8, close: u8) {
out.push(open);
let mut start = 0usize;
while let Some(found_rel) = find_close_quote_target(&data[start..], open, close) {
let found = start + found_rel;
if found > start {
out.extend_from_slice(&data[start..found]);
}
let byte = data[found];
out.push(byte);
out.push(byte);
start = found + 1;
}
out.extend_from_slice(&data[start..]);
out.push(close);
}
#[inline]
fn find_close_quote_or_backslash_target(haystack: &[u8], open: u8, close: u8) -> Option<usize> {
if open == close && close == b'\\' {
return memchr::memchr(b'\\', haystack);
}
let quote = if open == close { open } else { close };
if quote == b'\\' {
return memchr::memchr(b'\\', haystack);
}
memchr2(quote, b'\\', haystack)
}
fn write_quoted_backslash(data: &[u8], out: &mut Vec<u8>, open: u8, close: u8) {
out.push(open);
let mut start = 0usize;
while let Some(found_rel) = find_close_quote_or_backslash_target(&data[start..], open, close) {
let found = start + found_rel;
if found > start {
out.extend_from_slice(&data[start..found]);
}
out.push(b'\\');
out.push(data[found]);
start = found + 1;
}
out.extend_from_slice(&data[start..]);
out.push(close);
}
#[inline]
fn find_backslash_escape_target(haystack: &[u8], delimiter: u8) -> Option<usize> {
if delimiter == b'\n' || delimiter == b'\r' || delimiter == b'\\' {
return memchr3(b'\n', b'\r', b'\\', haystack);
}
match (memchr3(b'\n', b'\r', b'\\', haystack), memchr::memchr(delimiter, haystack)) {
(Some(left), Some(right)) => Some(left.min(right)),
(Some(left), None) => Some(left),
(None, Some(right)) => Some(right),
(None, None) => None,
}
}
fn write_backslash_escaped(data: &[u8], out: &mut Vec<u8>, delimiter: u8) {
let mut start = 0usize;
while let Some(found_rel) = find_backslash_escape_target(&data[start..], delimiter) {
let found = start + found_rel;
if found > start {
out.extend_from_slice(&data[start..found]);
}
match data[found] {
b'\n' => out.extend_from_slice(br"\n"),
b'\r' => out.extend_from_slice(br"\r"),
b'\\' => out.extend_from_slice(br"\\"),
byte if byte == delimiter => {
out.push(b'\\');
out.push(delimiter);
}
_ => unreachable!("find_backslash_escape_target returned non-special byte"),
}
start = found + 1;
}
out.extend_from_slice(&data[start..]);
}
#[inline]
fn find_replace_target(haystack: &[u8], delimiter: u8) -> Option<usize> {
if delimiter == b'\n' || delimiter == b'\r' {
memchr2(b'\n', b'\r', haystack)
} else {
memchr3(delimiter, b'\n', b'\r', haystack)
}
}
#[inline]
fn contains_quote_byte(data: &[u8], open: u8, close: u8) -> bool {
find_close_quote_target(data, open, close).is_some()
}
#[inline]
fn contains_quote_or_backslash(data: &[u8], open: u8, close: u8) -> bool {
find_close_quote_or_backslash_target(data, open, close).is_some()
}
fn write_replaced_or_deleted(
data: &[u8],
out: &mut Vec<u8>,
delimiter: u8,
replacement: Option<u8>,
) {
let Some(first) = find_replace_target(data, delimiter) else {
out.extend_from_slice(data);
return;
};
if first > 0 {
out.extend_from_slice(&data[..first]);
}
if let Some(byte) = replacement {
out.push(byte);
}
let mut start = first + 1;
while let Some(found_rel) = find_replace_target(&data[start..], delimiter) {
let found = start + found_rel;
if found > start {
out.extend_from_slice(&data[start..found]);
}
if let Some(byte) = replacement {
out.push(byte);
}
start = found + 1;
}
out.extend_from_slice(&data[start..]);
}
fn input_quote_pair(quotes: &Quotes) -> Option<(u8, u8)> {
match quotes {
Quotes::None => None,
Quotes::Single(open, close) => Some((*open, *close)),
Quotes::Multi(pairs) => pairs.first().copied(),
}
}
const fn derive_delimiter(input: &input_file::Config, spec: &Spec) -> u8 {
if let Some(delimiter) = spec.delimiter {
return delimiter;
}
match input.column_config.delimiter {
input::Delimiter::Char(delimiter) => delimiter,
_ => b'\t',
}
}
fn derive_escape(input: &input_file::Config, spec: &Spec, output_delimiter: u8) -> Escape {
if let Some(escape) = spec.escape {
return escape;
}
match (input.column_config.backslash, input_quote_pair(&input.column_config.quotes)) {
(BackslashMode::Off, None) => {
let replacement = if output_delimiter == b' ' { b'.' } else { b' ' };
Escape::Replace(replacement)
}
(BackslashMode::Off, Some((open, close))) => Escape::QuoteDoubled(open, close),
(BackslashMode::On, None) => Escape::Backslash,
(BackslashMode::On, Some((open, close))) => Escape::QuoteBackslash(open, close),
}
}
const fn derive_header(input: &input_file::Config, spec: &Spec) -> Header {
if let Some(header) = spec.header {
return header;
}
if let Some(saw_header) = input.saw_header {
return match saw_header {
input_file::SawHeader::Yes | input_file::SawHeader::Cdx => Header::Yes,
input_file::SawHeader::No | input_file::SawHeader::Empty => Header::No,
};
}
match input.header {
input_file::Header::Yes => Header::Yes,
input_file::Header::No => Header::No,
}
}
impl Spec {
#[must_use]
pub const fn new() -> Self {
Self { delimiter: None, escape: None, header: None, eol: None }
}
#[must_use]
pub const fn tsv() -> Self {
Self {
delimiter: Some(b'\t'),
escape: Some(Escape::Backslash),
header: Some(Header::Yes),
eol: None,
}
}
#[must_use]
pub const fn csv() -> Self {
Self {
delimiter: Some(b','),
escape: Some(Escape::QuoteDoubled(b'"', b'"')),
header: Some(Header::Yes),
eol: Some(Eol::Plain(ReadResult::CrLf)),
}
}
#[must_use]
pub const fn any() -> Self {
Self::new()
}
#[must_use]
pub fn eol(&self, input: ReadResult) -> &[u8] {
match self.eol {
None => input.as_bytes(),
Some(ref eol) => eol.as_bytes(),
}
}
pub fn from_spec(spec: &str) -> Result<Self> {
let trimmed = spec.trim();
if trimmed.is_empty() {
return Ok(Self::any());
}
let mut parts = trimmed.split(',').map(str::trim);
let first = parts.next().ok_or_else(|| anyhow::anyhow!("missing output spec"))?;
let mut out;
let mut escape_kind;
let mut quote_pair;
let mut replace_value;
let mut seen_keys: HashSet<&'static str> = HashSet::new();
let mut quote_overridden = false;
let mut replace_overridden = false;
if first.contains('=') {
out = Self::any();
escape_kind = None;
quote_pair = None;
replace_value = None;
} else {
let base_norm = normalize_token(first);
out = match base_norm.as_str() {
"csv" => Self::csv(),
"tsv" => Self::tsv(),
_ => {
return Err(anyhow::anyhow!(format!(
"unknown output spec base `{first}`; expected one of: csv, tsv"
)));
}
};
let (kind, pair, replace) =
escape_parts(out.escape.expect("Spec::csv and Spec::tsv always set escape"));
escape_kind = Some(kind);
quote_pair = Some(pair);
replace_value = replace;
}
let mut apply_override = |part: &str| -> Result<()> {
if part.is_empty() {
return Err(anyhow::anyhow!("empty override segment in output spec"));
}
let (key_raw, value_raw) = part.split_once('=').ok_or_else(|| {
anyhow::anyhow!(format!("output spec override `{part}` must be in key=value form"))
})?;
let key_norm = normalize_token(key_raw);
let key = canonical_output_key(&key_norm).ok_or_else(|| {
anyhow::anyhow!(format!(
"unknown output spec key `{key_raw}`; expected one of: d|delimiter, e|escape, h|header, q|quotes, x|replace"
))
})?;
if !seen_keys.insert(key) {
return Err(anyhow::anyhow!(format!("duplicate output spec key `{key_raw}`")));
}
match key {
"delimiter" => {
out.delimiter = Some(parse_byte_token(value_raw, "delimiter")?);
}
"escape" => {
escape_kind = Some(parse_escape_kind(value_raw)?);
}
"header" => {
out.header = Some(parse_header_mode(value_raw)?);
}
"quotes" => {
quote_pair = Some(parse_output_quote_pair(value_raw)?);
quote_overridden = true;
}
"replace" => {
replace_value = Some(parse_byte_token(value_raw, "replace")?);
replace_overridden = true;
}
_ => unreachable!("all output keys are canonicalized above"),
}
Ok(())
};
if first.contains('=') {
apply_override(first)?;
}
for part in parts {
apply_override(part)?;
}
if quote_overridden && !escape_kind.is_some_and(escape_kind_uses_quotes) {
return Err(anyhow::anyhow!(
"quotes override requires a quote-based escape mode via base or `e=`",
));
}
if replace_overridden && !matches!(escape_kind, Some(EscapeKind::Replace)) {
return Err(anyhow::anyhow!(
"replace override requires `escape=replace` via base or `e=`",
));
}
if let Some(kind) = escape_kind {
out.escape = Some(match kind {
EscapeKind::QuoteDoubled => {
let (open, close) = quote_pair.unwrap_or((b'"', b'"'));
Escape::QuoteDoubled(open, close)
}
EscapeKind::AlwaysQuoteDoubled => {
let (open, close) = quote_pair.unwrap_or((b'"', b'"'));
Escape::AlwaysQuoteDoubled(open, close)
}
EscapeKind::QuoteBackslash => {
let (open, close) = quote_pair.unwrap_or((b'"', b'"'));
Escape::QuoteBackslash(open, close)
}
EscapeKind::AlwaysQuoteBackslash => {
let (open, close) = quote_pair.unwrap_or((b'"', b'"'));
Escape::AlwaysQuoteBackslash(open, close)
}
EscapeKind::Backslash => Escape::Backslash,
EscapeKind::Error => Escape::Error,
EscapeKind::Trust => Escape::Trust,
EscapeKind::Delete => Escape::Delete,
EscapeKind::Replace => {
let value = replace_value
.ok_or_else(|| anyhow::anyhow!("`escape=replace` requires `x=<byte>`"))?;
Escape::Replace(value)
}
});
}
Ok(out)
}
}
impl Config {
#[must_use]
pub const fn new(delimiter: u8, escape: Escape) -> Self {
Self { delimiter, escape, header: Header::Yes }
}
#[must_use]
pub const fn with_header(delimiter: u8, escape: Escape, header: Header) -> Self {
Self { delimiter, escape, header }
}
#[must_use]
pub const fn tsv() -> Self {
Self { delimiter: b'\t', escape: Escape::Backslash, header: Header::Yes }
}
#[must_use]
pub const fn csv() -> Self {
Self { delimiter: b',', escape: Escape::QuoteDoubled(b'"', b'"'), header: Header::Yes }
}
#[must_use]
pub fn from_input_and_spec(input: &input_file::Config, spec: &Spec) -> Self {
let delimiter = derive_delimiter(input, spec);
let escape = derive_escape(input, spec, delimiter);
let header = derive_header(input, spec);
Self { delimiter, escape, header }
}
#[must_use]
pub fn from_spec(input: &TextFile, spec: &Spec) -> Self {
Self::from_input_and_spec(input.config(), spec)
}
pub fn write_column(&self, data: &[u8], out: &mut Vec<u8>) -> Result<()> {
match self.escape {
Escape::Trust => {
out.extend_from_slice(data);
Ok(())
}
Escape::Error => {
if has_delimiter_or_newline(data, self.delimiter) {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"column contains delimiter or newline",
)
.into());
}
out.extend_from_slice(data);
Ok(())
}
Escape::Replace(replacement) => {
write_replaced_or_deleted(data, out, self.delimiter, Some(replacement));
Ok(())
}
Escape::Delete => {
write_replaced_or_deleted(data, out, self.delimiter, None);
Ok(())
}
Escape::Backslash => {
write_backslash_escaped(data, out, self.delimiter);
Ok(())
}
Escape::QuoteDoubled(open, close) => {
let needs_quotes = has_delimiter_or_newline(data, self.delimiter)
|| contains_quote_byte(data, open, close);
if needs_quotes {
write_quoted_doubled(data, out, open, close);
} else {
out.extend_from_slice(data);
}
Ok(())
}
Escape::AlwaysQuoteDoubled(open, close) => {
write_quoted_doubled(data, out, open, close);
Ok(())
}
Escape::QuoteBackslash(open, close) => {
let needs_quotes = has_delimiter_or_newline(data, self.delimiter)
|| contains_quote_or_backslash(data, open, close);
if needs_quotes {
write_quoted_backslash(data, out, open, close);
} else {
out.extend_from_slice(data);
}
Ok(())
}
Escape::AlwaysQuoteBackslash(open, close) => {
write_quoted_backslash(data, out, open, close);
Ok(())
}
}
}
}
impl FromStr for Spec {
type Err = anyhow::Error;
fn from_str(s: &str) -> Result<Self> {
Self::from_spec(s)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum EscapeKind {
QuoteDoubled,
AlwaysQuoteDoubled,
QuoteBackslash,
AlwaysQuoteBackslash,
Backslash,
Error,
Trust,
Replace,
Delete,
}
fn normalize_token(token: &str) -> String {
token
.trim()
.chars()
.filter(|ch| *ch != '_' && *ch != '-')
.flat_map(char::to_lowercase)
.collect()
}
fn canonical_output_key(normalized: &str) -> Option<&'static str> {
match normalized {
"d" | "delimiter" => Some("delimiter"),
"e" | "escape" => Some("escape"),
"h" | "header" => Some("header"),
"q" | "quotes" => Some("quotes"),
"x" | "replace" => Some("replace"),
_ => None,
}
}
fn parse_byte_token(value: &str, label: &str) -> Result<u8> {
if let Some(prefix) = value.get(..3)
&& prefix.eq_ignore_ascii_case("ch:")
{
let literal = &value[3..];
let bytes = literal.as_bytes();
if bytes.len() == 1 && bytes[0].is_ascii() {
return Ok(bytes[0]);
}
return Err(anyhow::anyhow!(format!(
"{label} value `{value}` must use exactly one ASCII character after `ch:`"
)));
}
let trimmed = value.trim();
match normalize_token(trimmed).as_str() {
"t" | "tab" => Ok(b'\t'),
"s" | "space" => Ok(b' '),
"n" | "lf" | "newline" => Ok(b'\n'),
"r" | "cr" => Ok(b'\r'),
"c" | "comma" => Ok(b','),
"p" | "pipe" => Ok(b'|'),
"sc" | "semi" | "semicolon" => Ok(b';'),
_ => {
let bytes = trimmed.as_bytes();
if bytes.len() == 1 && bytes[0].is_ascii() {
Ok(bytes[0])
} else {
Err(anyhow::anyhow!(format!(
"unknown {label} value `{value}`; expected one of: t|tab, s|space, n|lf|newline, r|cr, c|comma, p|pipe, sc|semi|semicolon, one ASCII character, or ch:<char>"
)))
}
}
}
}
fn parse_escape_kind(value: &str) -> Result<EscapeKind> {
match normalize_token(value).as_str() {
"qd" | "quotedoubled" => Ok(EscapeKind::QuoteDoubled),
"aqd" | "alwaysquotedoubled" => Ok(EscapeKind::AlwaysQuoteDoubled),
"qb" | "quotebackslash" => Ok(EscapeKind::QuoteBackslash),
"aqb" | "alwaysquotebackslash" => Ok(EscapeKind::AlwaysQuoteBackslash),
"backslash" => Ok(EscapeKind::Backslash),
"error" => Ok(EscapeKind::Error),
"trust" => Ok(EscapeKind::Trust),
"replace" => Ok(EscapeKind::Replace),
"delete" => Ok(EscapeKind::Delete),
_ => Err(anyhow::anyhow!(format!(
"unknown escape value `{value}`; expected one of: qd|quote_doubled, aqd|always_quote_doubled, qb|quote_backslash, aqb|always_quote_backslash, backslash, error, trust, replace, delete"
))),
}
}
fn parse_header_mode(value: &str) -> Result<Header> {
match normalize_token(value).as_str() {
"y" | "yes" | "true" | "on" | "1" => Ok(Header::Yes),
"n" | "no" | "false" | "off" | "0" => Ok(Header::No),
_ => Err(anyhow::anyhow!(format!(
"unknown header value `{value}`; expected one of: y|yes|true|on|1, n|no|false|off|0"
))),
}
}
fn parse_output_quote_pair(value: &str) -> Result<(u8, u8)> {
match normalize_token(value).as_str() {
"ang" => return Ok((b'<', b'>')),
"brk" => return Ok((b'[', b']')),
"dq" => return Ok((b'"', b'"')),
"sq" => return Ok((b'\'', b'\'')),
_ => {}
}
let trimmed = value.trim();
let bytes = trimmed.as_bytes();
match bytes {
[single] if single.is_ascii() => Ok((*single, *single)),
[open, close] if open.is_ascii() && close.is_ascii() => Ok((*open, *close)),
_ => Err(anyhow::anyhow!(format!(
"quotes value `{value}` must be one ASCII character, two ASCII characters, or one of `ang`, `brk`, `dq`, `sq`"
))),
}
}
const fn escape_kind_uses_quotes(kind: EscapeKind) -> bool {
matches!(
kind,
EscapeKind::QuoteDoubled
| EscapeKind::AlwaysQuoteDoubled
| EscapeKind::QuoteBackslash
| EscapeKind::AlwaysQuoteBackslash
)
}
const fn escape_parts(escape: Escape) -> (EscapeKind, (u8, u8), Option<u8>) {
match escape {
Escape::QuoteDoubled(open, close) => (EscapeKind::QuoteDoubled, (open, close), None),
Escape::AlwaysQuoteDoubled(open, close) => {
(EscapeKind::AlwaysQuoteDoubled, (open, close), None)
}
Escape::QuoteBackslash(open, close) => (EscapeKind::QuoteBackslash, (open, close), None),
Escape::AlwaysQuoteBackslash(open, close) => {
(EscapeKind::AlwaysQuoteBackslash, (open, close), None)
}
Escape::Backslash => (EscapeKind::Backslash, (b'"', b'"'), None),
Escape::Error => (EscapeKind::Error, (b'"', b'"'), None),
Escape::Trust => (EscapeKind::Trust, (b'"', b'"'), None),
Escape::Replace(byte) => (EscapeKind::Replace, (b'"', b'"'), Some(byte)),
Escape::Delete => (EscapeKind::Delete, (b'"', b'"'), None),
}
}
pub type OutputConfig = Config;
impl Write for LineWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.write_column_chunk(buf).map_err(io::Error::other)?;
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
impl LineWriter {
#[must_use]
pub fn new(config: Config, eol: &[u8]) -> Self {
Self::with_capacities(config, eol, 0, 0)
}
#[must_use]
pub fn with_capacities(
config: Config,
eol: &[u8],
record_capacity: usize,
column_capacity: usize,
) -> Self {
Self {
config,
eol: eol.into(),
record_buf: Vec::with_capacity(record_capacity),
column_buf: Vec::with_capacity(column_capacity),
column_open: false,
wrote_any_column: false,
}
}
pub fn from_spec(&mut self, input: &TextFile, spec: &Spec) {
self.config = Config::from_spec(input, spec);
self.eol = spec.eol(input.eol()).into();
}
pub fn write_cdx(&self, w: &mut dyn Write) -> Result<()> {
match self.config.header {
Header::No => {}
Header::Yes => {
self.write(w)?;
}
Header::Cdx => {
w.write_all(b" CDX")?;
w.write_all(&[self.config.delimiter])?;
self.write(w)?;
}
}
Ok(())
}
#[must_use]
pub const fn config(&self) -> &Config {
&self.config
}
pub const fn config_mut(&mut self) -> &mut Config {
&mut self.config
}
#[must_use]
pub fn eol(&self) -> &[u8] {
&self.eol
}
pub const fn eol_mut(&mut self) -> &mut Vec<u8> {
&mut self.eol
}
#[must_use]
pub fn record(&self) -> &[u8] {
&self.record_buf
}
pub const fn record_mut(&mut self) -> &mut Vec<u8> {
&mut self.record_buf
}
#[must_use]
pub const fn is_column_open(&self) -> bool {
self.column_open
}
fn begin_column_write(&mut self) -> Result<()> {
if self.column_open {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"cannot write column while chunked column is open",
)
.into());
}
if self.wrote_any_column {
self.record_buf.push(self.config.delimiter);
} else {
self.wrote_any_column = true;
}
debug_assert!(!self.column_open);
Ok(())
}
pub fn write_column(&mut self, data: &[u8]) -> Result<()> {
let original_len = self.record_buf.len();
let original_wrote_any_column = self.wrote_any_column;
self.begin_column_write()?;
match self.config.write_column(data, &mut self.record_buf) {
Ok(()) => Ok(()),
Err(err) => {
self.record_buf.truncate(original_len);
self.wrote_any_column = original_wrote_any_column;
Err(err)
}
}
}
pub fn write_columns<I, B>(&mut self, columns: I) -> Result<()>
where
I: IntoIterator<Item = B>,
B: AsRef<[u8]>,
{
for column in columns {
self.write_column(column.as_ref())?;
}
Ok(())
}
pub fn write_column_chunks<I, B>(&mut self, chunks: I) -> Result<()>
where
I: IntoIterator<Item = B>,
B: AsRef<[u8]>,
{
self.begin_column()?;
for chunk in chunks {
self.write_column_chunk(chunk.as_ref())?;
}
self.end_column()
}
pub fn begin_column(&mut self) -> Result<()> {
if self.column_open {
return Err(
io::Error::new(io::ErrorKind::InvalidInput, "chunked column already open").into()
);
}
self.column_buf.clear();
self.column_open = true;
Ok(())
}
pub fn write_column_chunk(&mut self, data: &[u8]) -> Result<()> {
if !self.column_open {
return Err(
io::Error::new(io::ErrorKind::InvalidInput, "no open chunked column").into()
);
}
self.column_buf.extend_from_slice(data);
Ok(())
}
pub fn end_column(&mut self) -> Result<()> {
if !self.column_open {
return Err(
io::Error::new(io::ErrorKind::InvalidInput, "no open chunked column").into()
);
}
self.column_open = false;
let data = std::mem::take(&mut self.column_buf);
let result = self.write_column(&data);
self.column_buf = data;
result
}
pub fn clear(&mut self) {
self.record_buf.clear();
self.column_buf.clear();
self.column_open = false;
self.wrote_any_column = false;
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.record_buf.is_empty()
}
#[must_use]
pub const fn is_open(&self) -> bool {
self.column_open
}
pub fn write(&self, w: &mut dyn Write) -> Result<()> {
w.write_all(&self.record_buf)?;
w.write_all(&self.eol)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
fn output(config: OutputConfig, data: &[u8]) -> Result<Vec<u8>> {
let mut out = Vec::new();
config.write_column(data, &mut out)?;
Ok(out)
}
#[derive(Debug, Clone, Default)]
struct SharedVecWriter {
data: Arc<Mutex<Vec<u8>>>,
}
impl SharedVecWriter {
fn snapshot(&self) -> Vec<u8> {
self.data.lock().unwrap().clone()
}
}
impl Write for SharedVecWriter {
fn write(&mut self, buf: &[u8]) -> std::result::Result<usize, io::Error> {
self.data.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::result::Result<(), io::Error> {
Ok(())
}
}
#[derive(Debug, Clone, Default)]
struct CountingCallsWriter {
data: Arc<Mutex<Vec<u8>>>,
calls: Arc<AtomicUsize>,
}
impl CountingCallsWriter {
fn snapshot(&self) -> Vec<u8> {
self.data.lock().unwrap().clone()
}
fn calls(&self) -> usize {
self.calls.load(Ordering::Relaxed)
}
}
impl Write for CountingCallsWriter {
fn write(&mut self, buf: &[u8]) -> std::result::Result<usize, io::Error> {
self.calls.fetch_add(1, Ordering::Relaxed);
self.data.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::result::Result<(), io::Error> {
Ok(())
}
}
#[test]
fn write_column_trust_writes_raw() {
let config = OutputConfig::with_header(b',', Escape::Trust, Header::No);
assert_eq!(output(config, b"a,b\nc").unwrap(), b"a,b\nc");
}
#[test]
fn write_column_error_rejects_delimiter_or_newline() {
let config = OutputConfig::with_header(b',', Escape::Error, Header::No);
assert!(output(config, b"a,b").is_err());
assert!(output(config, b"a\nb").is_err());
assert!(output(config, b"a\rb").is_err());
assert_eq!(output(config, b"abc").unwrap(), b"abc");
}
#[test]
fn write_column_replace_substitutes_delimiter_and_newlines() {
let config = OutputConfig::with_header(b',', Escape::Replace(b'?'), Header::No);
assert_eq!(output(config, b"a,b\nc\rd").unwrap(), b"a?b?c?d");
}
#[test]
fn write_column_delete_removes_delimiter_and_newlines() {
let config = OutputConfig::with_header(b',', Escape::Delete, Header::No);
assert_eq!(output(config, b"a,b\nc\rd").unwrap(), b"abcd");
}
#[test]
fn write_column_backslash_escapes_special_bytes() {
let config = OutputConfig::with_header(b',', Escape::Backslash, Header::No);
assert_eq!(output(config, b"a,b\nc\rd\\e").unwrap(), b"a\\,b\\nc\\rd\\\\e");
}
#[test]
fn write_column_quote_doubled_quotes_when_needed() {
let config = OutputConfig::with_header(b',', Escape::QuoteDoubled(b'"', b'"'), Header::No);
assert_eq!(output(config, b"a,b").unwrap(), b"\"a,b\"");
assert_eq!(output(config, b"ab").unwrap(), b"ab");
assert_eq!(output(config, b"a\"b").unwrap(), b"\"a\"\"b\"");
}
#[test]
fn write_column_always_quote_doubled_always_quotes() {
let config =
OutputConfig::with_header(b',', Escape::AlwaysQuoteDoubled(b'"', b'"'), Header::No);
assert_eq!(output(config, b"ab").unwrap(), b"\"ab\"");
assert_eq!(output(config, b"a\"b").unwrap(), b"\"a\"\"b\"");
}
#[test]
fn write_column_quote_doubled_asymmetric_only_doubles_close_quote() {
let config = OutputConfig::with_header(b',', Escape::QuoteDoubled(b'<', b'>'), Header::No);
assert_eq!(output(config, b"a<b").unwrap(), b"a<b");
assert_eq!(output(config, b"a>b<c").unwrap(), b"<a>>b<c>");
}
#[test]
fn write_column_always_quote_doubled_asymmetric_only_doubles_close_quote() {
let config =
OutputConfig::with_header(b',', Escape::AlwaysQuoteDoubled(b'<', b'>'), Header::No);
assert_eq!(output(config, b"a<b").unwrap(), b"<a<b>");
assert_eq!(output(config, b"a>b<c").unwrap(), b"<a>>b<c>");
}
#[test]
fn write_column_quote_backslash_quotes_when_needed() {
let config =
OutputConfig::with_header(b',', Escape::QuoteBackslash(b'"', b'"'), Header::No);
assert_eq!(output(config, b"a,b").unwrap(), b"\"a,b\"");
assert_eq!(output(config, b"ab").unwrap(), b"ab");
assert_eq!(output(config, b"a\"b\\c").unwrap(), b"\"a\\\"b\\\\c\"");
}
#[test]
fn write_column_quote_backslash_asymmetric_escapes_close_and_backslash_only() {
let config =
OutputConfig::with_header(b',', Escape::QuoteBackslash(b'<', b'>'), Header::No);
assert_eq!(output(config, b"a<b").unwrap(), b"a<b");
assert_eq!(output(config, b"a>b\\c").unwrap(), b"<a\\>b\\\\c>");
}
#[test]
fn write_column_always_quote_backslash_always_quotes() {
let config =
OutputConfig::with_header(b',', Escape::AlwaysQuoteBackslash(b'"', b'"'), Header::No);
assert_eq!(output(config, b"ab").unwrap(), b"\"ab\"");
assert_eq!(output(config, b"a\"b\\c").unwrap(), b"\"a\\\"b\\\\c\"");
}
#[test]
fn output_config_safe_constructors() {
assert_eq!(
OutputConfig::tsv(),
OutputConfig { delimiter: b'\t', escape: Escape::Backslash, header: Header::Yes }
);
assert_eq!(
OutputConfig::csv(),
OutputConfig {
delimiter: b',',
escape: Escape::QuoteDoubled(b'"', b'"'),
header: Header::Yes,
}
);
}
#[test]
fn output_spec_from_spec_parses_aliases() {
let parsed = Spec::from_spec("csv,d=t,e=aqb,q=ang").unwrap();
assert_eq!(
parsed,
Spec {
delimiter: Some(b'\t'),
escape: Some(Escape::AlwaysQuoteBackslash(b'<', b'>')),
header: Some(Header::Yes),
eol: Some(Eol::Plain(ReadResult::CrLf))
}
);
}
#[test]
fn output_spec_from_spec_quote_values() {
let one = Spec::from_spec("csv,e=qd,q='").unwrap();
assert_eq!(one.escape, Some(Escape::QuoteDoubled(b'\'', b'\'')));
assert_eq!(one.header, Some(Header::Yes));
let two = Spec::from_spec("csv,e=qb,q=[]").unwrap();
assert_eq!(two.escape, Some(Escape::QuoteBackslash(b'[', b']')));
assert_eq!(two.header, Some(Header::Yes));
}
#[test]
fn output_spec_from_spec_header_values() {
let yes = Spec::from_spec("csv,h=yes").unwrap();
assert_eq!(yes.header, Some(Header::Yes));
let no = Spec::from_spec("csv,header=off").unwrap();
assert_eq!(no.header, Some(Header::No));
}
#[test]
fn output_spec_from_spec_replace_requires_value() {
let err = Spec::from_spec("csv,e=replace").unwrap_err();
assert!(err.to_string().contains("requires"));
let parsed = Spec::from_spec("csv,e=replace,x=s").unwrap();
assert_eq!(parsed.escape, Some(Escape::Replace(b' ')));
}
#[test]
fn output_spec_from_spec_rejects_quotes_for_non_quote_escape() {
let err = Spec::from_spec("csv,e=delete,q=brk").unwrap_err();
assert!(err.to_string().contains("quote-based"));
}
#[test]
fn output_spec_from_str_works() {
let parsed: Spec = "tsv,e=trust,d=s".parse().unwrap();
assert_eq!(parsed.delimiter, Some(b' '));
assert_eq!(parsed.escape, Some(Escape::Trust));
}
#[test]
fn output_spec_from_spec_unknown_base_lists_expected_values() {
let err = Spec::from_spec("nope").unwrap_err();
let text = err.to_string();
assert!(text.contains("expected one of"));
assert!(text.contains("csv"));
assert!(text.contains("tsv"));
}
#[test]
fn output_spec_from_spec_omitted_base_allows_key_value_start() {
let parsed = Spec::from_spec("d=t,e=backslash").unwrap();
assert_eq!(parsed.delimiter, Some(b'\t'));
assert_eq!(parsed.escape, Some(Escape::Backslash));
assert_eq!(parsed.header, None);
}
#[test]
fn output_spec_from_spec_empty_is_any() {
assert_eq!(Spec::from_spec("").unwrap(), Spec::any());
assert_eq!(Spec::from_spec(" ").unwrap(), Spec::any());
}
#[test]
fn output_spec_from_spec_quotes_without_escape_is_error() {
let err = Spec::from_spec("q=dq").unwrap_err();
assert!(err.to_string().contains("quote-based"));
}
#[test]
fn output_spec_from_spec_replace_without_escape_is_error() {
let err = Spec::from_spec("x=s").unwrap_err();
assert!(err.to_string().contains("escape=replace"));
}
#[test]
fn output_spec_new_is_all_none() {
assert_eq!(Spec::new(), Spec { delimiter: None, escape: None, header: None, eol: None });
}
#[test]
fn output_spec_any_matches_new() {
assert_eq!(Spec::any(), Spec::new());
}
#[test]
fn output_spec_csv_tsv_match_config_defaults() {
assert_eq!(
Spec::csv(),
Spec {
delimiter: Some(Config::csv().delimiter),
escape: Some(Config::csv().escape),
header: Some(Config::csv().header),
eol: Some(Eol::Plain(ReadResult::CrLf))
}
);
assert_eq!(
Spec::tsv(),
Spec {
delimiter: Some(Config::tsv().delimiter),
escape: Some(Config::tsv().escape),
header: Some(Config::tsv().header),
eol: None,
}
);
}
#[test]
fn output_config_from_input_and_spec_csv_defaults() {
let input = input_file::Config::csv();
let output = Config::from_input_and_spec(&input, &Spec::new());
assert_eq!(output.delimiter, b',');
assert_eq!(output.escape, Escape::QuoteDoubled(b'"', b'"'));
assert_eq!(output.header, Header::Yes);
}
#[test]
fn output_config_from_input_and_spec_non_char_delimiter_fallback() {
let input = input_file::Config::whole();
let output = Config::from_input_and_spec(&input, &Spec::new());
assert_eq!(output.delimiter, b'\t');
assert_eq!(output.escape, Escape::Replace(b' '));
assert_eq!(output.header, Header::No); }
#[test]
fn output_config_from_input_and_spec_replace_dot_for_space_delimiter() {
let input = input_file::Config::from_delim(input::Delimiter::Char(b','));
let spec = Spec { delimiter: Some(b' '), escape: None, header: None, eol: None };
let output = Config::from_input_and_spec(&input, &spec);
assert_eq!(output.escape, Escape::Replace(b'.'));
}
#[test]
fn output_config_from_input_and_spec_multi_quotes_uses_first_pair() {
let mut input = input_file::Config::from_delim(input::Delimiter::Char(b','));
input.column_config.backslash = BackslashMode::On;
input.column_config.quotes = Quotes::Multi(vec![(b'<', b'>'), (b'[', b']')]);
let output = Config::from_input_and_spec(&input, &Spec::new());
assert_eq!(output.escape, Escape::QuoteBackslash(b'<', b'>'));
}
#[test]
fn output_config_from_input_and_spec_header_precedence() {
let mut input = input_file::Config::csv();
input.header = input_file::Header::No;
input.saw_header = Some(input_file::SawHeader::Cdx);
let output = Config::from_input_and_spec(&input, &Spec::new());
assert_eq!(output.header, Header::Yes);
input.saw_header = Some(input_file::SawHeader::No);
let output = Config::from_input_and_spec(&input, &Spec::new());
assert_eq!(output.header, Header::No);
}
#[test]
fn output_config_from_input_and_spec_overrides() {
let input = input_file::Config::csv();
let spec = Spec {
delimiter: Some(b'|'),
escape: Some(Escape::Delete),
header: Some(Header::No),
eol: None,
};
let output = Config::from_input_and_spec(&input, &spec);
assert_eq!(output, Config { delimiter: b'|', escape: Escape::Delete, header: Header::No });
}
#[test]
fn line_writer_write_record_writes_columns_and_eol() {
let config = OutputConfig::with_header(b',', Escape::QuoteDoubled(b'"', b'"'), Header::No);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_columns([&b"a,b"[..], &b"c"[..]]).unwrap();
assert_eq!(line_writer.record(), b"\"a,b\",c");
line_writer.clear();
line_writer.write_columns([&b"x"[..], &b"y"[..]]).unwrap();
assert_eq!(line_writer.record(), b"x,y");
}
#[test]
fn line_writer_with_capacities_uses_requested_capacity() {
let config = OutputConfig::with_header(b',', Escape::Trust, Header::No);
let line_writer = LineWriter::with_capacities(config, b"\n", 64, 32);
assert!(line_writer.record_buf.capacity() >= 64);
assert!(line_writer.column_buf.capacity() >= 32);
}
#[test]
fn line_writer_is_column_open_tracks_state() {
let config = OutputConfig::new(b',', Escape::Trust);
let mut line_writer = LineWriter::new(config, b"\n");
assert!(!line_writer.is_column_open());
line_writer.begin_column().unwrap();
assert!(line_writer.is_column_open());
line_writer.write_column_chunk(b"a").unwrap();
line_writer.end_column().unwrap();
assert!(!line_writer.is_column_open());
}
#[test]
fn line_writer_abort_record_discards_pending_data() {
let config = OutputConfig::new(b',', Escape::Trust);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_column(b"left").unwrap();
line_writer.begin_column().unwrap();
line_writer.write_column_chunk(b"right").unwrap();
line_writer.clear();
assert!(!line_writer.is_column_open());
line_writer.write_column(b"kept").unwrap();
assert_eq!(line_writer.record(), b"kept");
}
#[test]
fn line_writer_write_columns_then_finish_record() {
let mut sink = SharedVecWriter::default();
let sink_view = sink.clone();
let config = OutputConfig::with_header(b'\t', Escape::Trust, Header::No);
let mut line_writer = LineWriter::new(config, b"\r\n");
line_writer.write_columns([&b"aa"[..], &b"bb"[..], &b"cc"[..]]).unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"aa\tbb\tcc\r\n");
}
#[test]
fn line_writer_write_column_writes_tab_delimiter_and_lf_eol() {
let mut sink = SharedVecWriter::default();
let sink_view = sink.clone();
let config = OutputConfig::with_header(b'\t', Escape::Trust, Header::No);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_column(b"foo").unwrap();
line_writer.write_column(b"bar").unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"foo\tbar\n");
}
#[test]
fn line_writer_chunked_column_mode() {
let mut sink = SharedVecWriter::default();
let sink_view = sink.clone();
let config =
OutputConfig::with_header(b',', Escape::QuoteBackslash(b'"', b'"'), Header::No);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_column(b"left").unwrap();
line_writer.begin_column().unwrap();
line_writer.write_column_chunk(b"a\"").unwrap();
line_writer.write_column_chunk(b"b\\c").unwrap();
line_writer.end_column().unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"left,\"a\\\"b\\\\c\"\n");
}
#[test]
fn line_writer_write_column_chunks_convenience_api() {
let mut sink = SharedVecWriter::default();
let sink_view = sink.clone();
let config =
OutputConfig::with_header(b',', Escape::QuoteBackslash(b'"', b'"'), Header::No);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_column(b"left").unwrap();
line_writer.write_column_chunks([&b"a\""[..], &b"b\\c"[..]]).unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"left,\"a\\\"b\\\\c\"\n");
}
#[test]
fn line_writer_chunked_column_chunked_and_bulk_api() {
let mut sink = SharedVecWriter::default();
let sink_view = sink.clone();
let config =
OutputConfig::with_header(b',', Escape::QuoteBackslash(b'"', b'"'), Header::No);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_column(b"left").unwrap();
line_writer.begin_column().unwrap();
line_writer.write_column_chunk(b"a\"").unwrap();
line_writer.write_column_chunk(b"b\\c").unwrap();
line_writer.end_column().unwrap();
line_writer.write_column_chunks([&b"x"[..], &b"y"[..]]).unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"left,\"a\\\"b\\\\c\",xy\n");
}
#[test]
fn line_writer_finish_record_uses_single_write_call() {
let mut sink = CountingCallsWriter::default();
let sink_view = sink.clone();
let config = OutputConfig::with_header(b'\t', Escape::Trust, Header::No);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_column(b"foo").unwrap();
line_writer.write_column(b"bar").unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"foo\tbar\n");
assert_eq!(sink_view.calls(), 2);
}
#[test]
fn line_writer_write_column_error_rolls_back_state() {
let mut sink = SharedVecWriter::default();
let sink_view = sink.clone();
let config = OutputConfig::with_header(b',', Escape::Error, Header::No);
let mut line_writer = LineWriter::new(config, b"\n");
line_writer.write_column(b"left").unwrap();
assert!(line_writer.write_column(b"bad,value").is_err());
line_writer.write_column(b"right").unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"left,right\n");
}
#[test]
fn line_writer_finish_empty_record_writes_only_eol() {
let mut sink = SharedVecWriter::default();
let sink_view = sink.clone();
let config = OutputConfig::new(b',', Escape::Trust);
let mut line_writer = LineWriter::new(config, b"\r\n");
line_writer.write(&mut sink).unwrap();
line_writer.write_column(b"a").unwrap();
line_writer.write(&mut sink).unwrap();
assert_eq!(sink_view.snapshot(), b"\r\na\r\n");
}
}