diff --git a/Lib/test/test_csv.py b/Lib/test/test_csv.py index 0e1f020d3..494cce50a 100644 --- a/Lib/test/test_csv.py +++ b/Lib/test/test_csv.py @@ -854,7 +854,6 @@ class EscapedExcel(csv.excel): class TestEscapedExcel(TestCsvBase): dialect = EscapedExcel() - @unittest.expectedFailure # TODO: RUSTPYTHON def test_escape_fieldsep(self): self.writerAssertEqual([['abc,def']], 'abc\\,def\r\n') diff --git a/crates/stdlib/src/csv.rs b/crates/stdlib/src/csv.rs index bd80cf7ec..aaffab182 100644 --- a/crates/stdlib/src/csv.rs +++ b/crates/stdlib/src/csv.rs @@ -910,12 +910,8 @@ mod _csv { writer = writer.delimiter(t); } - if let Some(t) = self.quotechar { - if let Some(u) = t { - writer = writer.quote(u); - } else { - todo!() - } + if let Some(Some(t)) = self.quotechar { + writer = writer.quote(t); } if let Some(t) = self.doublequote { @@ -1148,6 +1144,24 @@ mod _csv { Ok(()) } + fn write_unquoted_field( + output: &mut Vec, + data: &[u8], + dialect: PyDialect, + vm: &VirtualMachine, + ) -> PyResult<()> { + for &byte in data { + if field_needs_escape(byte, dialect) { + let escapechar = dialect + .escapechar + .ok_or_else(|| new_csv_error(vm, "need to escape, but no escapechar set"))?; + output.push(escapechar); + } + output.push(byte); + } + Ok(()) + } + fn field_needs_quotes(data: &[u8], dialect: PyDialect) -> bool { data.iter().any(|&byte| { byte == dialect.delimiter @@ -1157,6 +1171,14 @@ mod _csv { }) } + fn field_needs_escape(byte: u8, dialect: PyDialect) -> bool { + byte == dialect.delimiter + || dialect.quotechar == Some(byte) + || dialect.escapechar == Some(byte) + || matches!(byte, b'\r' | b'\n') + || matches!(dialect.lineterminator, Terminator::Any(t) if byte == t) + } + fn write_lineterminator(output: &mut Vec, terminator: Terminator) { match terminator { Terminator::CRLF => output.extend_from_slice(b"\r\n"), @@ -1222,13 +1244,61 @@ mod _csv { self.write.call((s,), vm) } + fn writerow_quote_none(&self, row: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let _state = self.state.lock(); + + let row: ArgIterable = ArgIterable::try_from_object(vm, row.clone()).map_err(|_e| { + new_csv_error( + vm, + format!("'{}' object is not iterable", row.class().name()), + ) + })?; + + let fields = row.iter(vm)?.collect::>>()?; + let single_field = fields.len() == 1; + let mut output = Vec::new(); + + for (index, field) in fields.into_iter().enumerate() { + if index > 0 { + output.push(self.dialect.delimiter); + } + + let stringified; + let data: &[u8] = match_class!(match field { + ref s @ PyStr => s.as_bytes(), + crate::builtins::PyNone => b"", + ref obj => { + stringified = obj.str(vm)?; + stringified.as_bytes() + } + }); + + if single_field && data.is_empty() { + return Err(new_csv_error( + vm, + "single empty field record must be quoted", + )); + } + + write_unquoted_field(&mut output, data, self.dialect, vm)?; + } + + write_lineterminator(&mut output, self.dialect.lineterminator); + + let s = core::str::from_utf8(&output) + .map_err(|_| vm.new_unicode_decode_error("csv not utf8"))?; + + self.write.call((s,), vm) + } + #[pymethod] fn writerow(&self, row: PyObjectRef, vm: &VirtualMachine) -> PyResult { - if matches!( - self.dialect.quoting, - QuoteStyle::Strings | QuoteStyle::Notnull - ) { - return self.writerow_quoted_strings(row, vm); + match self.dialect.quoting { + QuoteStyle::None => return self.writerow_quote_none(row, vm), + QuoteStyle::Strings | QuoteStyle::Notnull => { + return self.writerow_quoted_strings(row, vm); + } + _ => {} } let mut state = self.state.lock(); diff --git a/extra_tests/snippets/stdlib_csv.py b/extra_tests/snippets/stdlib_csv.py index aa7b41223..b9c741cbb 100644 --- a/extra_tests/snippets/stdlib_csv.py +++ b/extra_tests/snippets/stdlib_csv.py @@ -72,3 +72,65 @@ def test_quote_strings_and_notnull_writer(): test_quote_strings_and_notnull_writer() + + +def test_quote_none_writer_without_quotechar(): + no_quotechar_buf = io.StringIO() + csv.writer( + no_quotechar_buf, + quoting=csv.QUOTE_NONE, + quotechar=None, + escapechar="\\", + ).writerow(["a,b", 'x"y']) + assert no_quotechar_buf.getvalue() == 'a\\,b,x"y\r\n' + + default_quotechar_buf = io.StringIO() + csv.writer( + default_quotechar_buf, + quoting=csv.QUOTE_NONE, + escapechar="\\", + ).writerow(["a,b", 'x"y']) + assert default_quotechar_buf.getvalue() == 'a\\,b,x\\"y\r\n' + + escapechar_buf = io.StringIO() + csv.writer( + escapechar_buf, + quoting=csv.QUOTE_NONE, + quotechar=None, + escapechar="\\", + ).writerow(["a\\b"]) + assert escapechar_buf.getvalue() == "a\\\\b\r\n" + + linebreak_buf = io.StringIO() + csv.writer( + linebreak_buf, + quoting=csv.QUOTE_NONE, + quotechar=None, + escapechar="\\", + ).writerow(["a\rb", "c\nd"]) + assert linebreak_buf.getvalue() == "a\\\rb,c\\\nd\r\n" + + with assert_raises(csv.Error): + csv.writer(io.StringIO(), quoting=csv.QUOTE_NONE, quotechar=None).writerow( + ["a,b"] + ) + + with assert_raises(csv.Error): + csv.writer( + io.StringIO(), + quoting=csv.QUOTE_NONE, + quotechar=None, + escapechar="\\", + ).writerow([None]) + + two_empty_buf = io.StringIO() + csv.writer( + two_empty_buf, + quoting=csv.QUOTE_NONE, + quotechar=None, + escapechar="\\", + ).writerow([None, ""]) + assert two_empty_buf.getvalue() == ",\r\n" + + +test_quote_none_writer_without_quotechar()