Add csv.writer

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__
 

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<Self, ArgumentError> {
         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<ReadState>,
 }
 
-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<u8>,
+    writer: csv_core::Writer,
+}
+
+#[pyclass(module = "_csv", name = "writer")]
+struct Writer {
+    write: PyObjectRef,
+    state: PyMutex<WriteState>,
+}
+
+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<Writer> {
+    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<T: num_traits::PrimInt>(buf: &mut Vec<T>) {
+    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),