diff --git a/Lib/csv.py b/Lib/csv.py index dd896f28c..2f38bb1a1 100644 --- a/Lib/csv.py +++ b/Lib/csv.py @@ -4,7 +4,7 @@ csv.py - read/write/investigate CSV files """ import re -from _csv import Error, reader, \ +from _csv import Error, writer, reader, \ QUOTE_MINIMAL, QUOTE_ALL, QUOTE_NONNUMERIC, QUOTE_NONE, \ __doc__ diff --git a/vm/src/stdlib/csv.rs b/vm/src/stdlib/csv.rs index 6e06724b6..fd297ab51 100644 --- a/vm/src/stdlib/csv.rs +++ b/vm/src/stdlib/csv.rs @@ -1,5 +1,5 @@ use itertools::{self, Itertools}; -use std::fmt::{self, Debug, Formatter}; +use std::fmt; use crate::builtins::pystr::{PyStr, PyStrRef}; use crate::builtins::pytype::PyTypeRef; @@ -7,8 +7,8 @@ use crate::common::lock::PyMutex; use crate::function::{ArgumentError, FromArgs, FuncArgs}; use crate::iterator; use crate::pyobject::{ - BorrowValue, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, StaticType, TryFromObject, - TypeProtocol, + BorrowValue, PyClassImpl, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, StaticType, + TryFromObject, TypeProtocol, }; use crate::slots::PyIter; use crate::types::create_simple_type; @@ -22,12 +22,12 @@ pub enum QuoteStyle { QuoteNone, } -struct ReaderOption { +struct FormatOptions { delimiter: u8, quotechar: u8, } -impl FromArgs for ReaderOption { +impl FromArgs for FormatOptions { fn from_args(vm: &VirtualMachine, args: &mut FuncArgs) -> Result { let delimiter = if let Some(delimiter) = args.kwargs.remove("delimiter") { PyStrRef::try_from_object(vm, delimiter)? @@ -55,18 +55,26 @@ impl FromArgs for ReaderOption { b'"' }; - Ok(ReaderOption { + Ok(FormatOptions { delimiter, quotechar, }) } } -impl ReaderOption { +impl FormatOptions { fn to_reader(&self) -> csv_core::Reader { csv_core::ReaderBuilder::new() .delimiter(self.delimiter) .quote(self.quotechar) + .terminator(csv_core::Terminator::CRLF) + .build() + } + fn to_writer(&self) -> csv_core::Writer { + csv_core::WriterBuilder::new() + .delimiter(self.delimiter) + .quote(self.quotechar) + .terminator(csv_core::Terminator::CRLF) .build() } } @@ -77,14 +85,14 @@ struct ReadState { reader: csv_core::Reader, } -#[pyclass(module = "csv", name = "Reader")] +#[pyclass(module = "_csv", name = "reader")] struct Reader { iter: PyObjectRef, state: PyMutex, } -impl Debug for Reader { - fn fmt(&self, f: &mut Formatter) -> fmt::Result { +impl fmt::Debug for Reader { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "_csv.reader") } } @@ -130,14 +138,8 @@ impl PyIter for Reader { output_ends_offset += nends; match res { csv_core::ReadRecordResult::InputEmpty => {} - csv_core::ReadRecordResult::OutputFull => { - let new_size = buffer.len() * 2; - buffer.resize(new_size, 0u8); - } - csv_core::ReadRecordResult::OutputEndsFull => { - let new_size = output_ends.len() * 2; - output_ends.resize(new_size, 0); - } + csv_core::ReadRecordResult::OutputFull => resize_buf(buffer), + csv_core::ReadRecordResult::OutputEndsFull => resize_buf(output_ends), csv_core::ReadRecordResult::Record => break, csv_core::ReadRecordResult::End => return Err(vm.new_stop_iteration()), } @@ -169,7 +171,7 @@ impl PyIter for Reader { fn _csv_reader( iter: PyObjectRef, - options: ReaderOption, + options: FormatOptions, // TODO: handle quote style, etc _rest: FuncArgs, vm: &VirtualMachine, @@ -185,16 +187,134 @@ fn _csv_reader( }) } +struct WriteState { + buffer: Vec, + writer: csv_core::Writer, +} + +#[pyclass(module = "_csv", name = "writer")] +struct Writer { + write: PyObjectRef, + state: PyMutex, +} + +impl fmt::Debug for Writer { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "_csv.writer") + } +} + +impl PyValue for Writer { + fn class(_vm: &VirtualMachine) -> &PyTypeRef { + Self::static_type() + } +} + +#[pyimpl] +impl Writer { + #[pymethod] + fn writerow(&self, row: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let mut state = self.state.lock(); + let WriteState { buffer, writer } = &mut *state; + + let mut buffer_offset = 0; + + macro_rules! handle_res { + ($x:expr) => {{ + let (res, nwritten) = $x; + buffer_offset += nwritten; + match res { + csv_core::WriteResult::InputEmpty => break, + csv_core::WriteResult::OutputFull => resize_buf(buffer), + } + }}; + } + + let row = PyIterable::try_from_object(vm, row)?; + for field in row.iter(vm)? { + let field: PyObjectRef = field?; + let stringified; + let data: &[u8] = match_class!(match field { + ref s @ PyStr => s.borrow_value().as_bytes(), + crate::builtins::PyNone => b"", + ref obj => { + stringified = vm.to_str(obj)?; + stringified.borrow_value().as_bytes() + } + }); + + let mut input_offset = 0; + + loop { + let (res, nread, nwritten) = + writer.field(&data[input_offset..], &mut buffer[buffer_offset..]); + input_offset += nread; + handle_res!((res, nwritten)); + } + + loop { + handle_res!(writer.delimiter(&mut buffer[buffer_offset..])); + } + } + + loop { + handle_res!(writer.terminator(&mut buffer[buffer_offset..])); + } + + let s = std::str::from_utf8(&buffer[..buffer_offset]) + .map_err(|_| vm.new_unicode_decode_error("csv not utf8".to_owned()))?; + + vm.invoke(&self.write, (s.to_owned(),)) + } + + #[pymethod] + fn writerows(&self, rows: PyIterable, vm: &VirtualMachine) -> PyResult<()> { + for row in rows.iter(vm)? { + self.writerow(row?, vm)?; + } + Ok(()) + } +} + +fn _csv_writer( + file: PyObjectRef, + options: FormatOptions, + // TODO: handle quote style, etc + _rest: FuncArgs, + vm: &VirtualMachine, +) -> PyResult { + let write = match vm.get_attribute_opt(file.clone(), "write")? { + Some(write_meth) => write_meth, + None if vm.is_callable(&file) => file, + None => return Err(vm.new_type_error("argument 1 must have a \"write\" method".to_owned())), + }; + + Ok(Writer { + write, + state: PyMutex::new(WriteState { + buffer: vec![0; 1024], + writer: options.to_writer(), + }), + }) +} + +#[inline] +fn resize_buf(buf: &mut Vec) { + let new_size = buf.len() * 2; + buf.resize(new_size, T::zero()); +} + pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; - let reader_type = Reader::make_class(ctx); + Reader::make_class(ctx); + Writer::make_class(ctx); let error = create_simple_type("Error", &ctx.exceptions.exception_type); py_module!(vm, "_csv", { "reader" => named_function!(ctx, _csv, reader), - "Reader" => reader_type, + "writer" => named_function!(ctx, _csv, writer), "Error" => error, // constants "QUOTE_MINIMAL" => ctx.new_int(QuoteStyle::QuoteMinimal as i32),