use csv as rust_csv; use itertools::{self, Itertools}; use std::fmt::{self, Debug, Formatter}; use crate::builtins::pystr::{self, PyStr}; use crate::builtins::pytype::PyTypeRef; use crate::common::lock::PyRwLock; use crate::function::FuncArgs; use crate::pyobject::{ BorrowValue, IntoPyObject, PyClassImpl, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, StaticType, TryFromObject, TypeProtocol, }; use crate::types::create_simple_type; use crate::VirtualMachine; #[repr(i32)] pub enum QuoteStyle { QuoteMinimal, QuoteAll, QuoteNonnumeric, QuoteNone, } struct ReaderOption { delimiter: u8, quotechar: u8, } impl ReaderOption { fn new(args: FuncArgs, vm: &VirtualMachine) -> PyResult { let delimiter = if let Some(delimiter) = args.get_optional_kwarg("delimiter") { *pystr::borrow_value(&delimiter) .as_bytes() .iter() .exactly_one() .map_err(|_| { let msg = r#""delimiter" must be a 1-character string"#; vm.new_type_error(msg.to_owned()) })? } else { b',' }; let quotechar = if let Some(quotechar) = args.get_optional_kwarg("quotechar") { *pystr::borrow_value("echar) .as_bytes() .iter() .exactly_one() .map_err(|_| { let msg = r#""quotechar" must be a 1-character string"#; vm.new_type_error(msg.to_owned()) })? } else { b'"' }; Ok(ReaderOption { delimiter, quotechar, }) } } pub fn build_reader( iterable: PyIterable, args: FuncArgs, vm: &VirtualMachine, ) -> PyResult { let config = ReaderOption::new(args, vm)?; Ok(Reader::new(iterable, config).into_object(vm)) } fn into_strings(iterable: &PyIterable, vm: &VirtualMachine) -> PyResult> { iterable .iter(vm)? .map(|py_obj_ref| { match_class!(match py_obj_ref? { py_str @ PyStr => Ok(py_str.borrow_value().trim().to_owned()), obj => { let msg = format!( "iterator should return strings, not {} (did you open the file in text mode?)", obj.class().name ); Err(vm.new_type_error(msg)) } }) }) .collect::>>() } type MemIO = std::io::Cursor>; #[allow(dead_code)] enum ReadState { PyIter(PyIterable, ReaderOption), CsvIter(rust_csv::StringRecordsIntoIter), } impl ReadState { fn new(iter: PyIterable, config: ReaderOption) -> Self { ReadState::PyIter(iter, config) } fn cast_to_reader(&mut self, vm: &VirtualMachine) -> PyResult<()> { if let ReadState::PyIter(ref iterable, ref config) = self { let lines = into_strings(iterable, vm)?; let contents = itertools::join(lines, "\n"); let bytes = Vec::from(contents.as_bytes()); let reader = MemIO::new(bytes); let csv_iter = rust_csv::ReaderBuilder::new() .delimiter(config.delimiter) .quote(config.quotechar) .has_headers(false) .from_reader(reader) .into_records(); *self = ReadState::CsvIter(csv_iter); } Ok(()) } } #[pyclass(module = "csv", name = "Reader")] struct Reader { state: PyRwLock, } impl Debug for Reader { fn fmt(&self, f: &mut Formatter) -> fmt::Result { write!(f, "_csv.reader") } } impl PyValue for Reader { fn class(_vm: &VirtualMachine) -> &PyTypeRef { Self::static_type() } } impl Reader { fn new(iter: PyIterable, config: ReaderOption) -> Self { let state = PyRwLock::new(ReadState::new(iter, config)); Reader { state } } } #[pyimpl] impl Reader { #[pymethod(name = "__iter__")] fn iter(this: PyRef, vm: &VirtualMachine) -> PyResult { this.state.write().cast_to_reader(vm)?; Ok(this.into_pyobject(vm)) } #[pymethod(name = "__next__")] fn next(&self, vm: &VirtualMachine) -> PyResult { let mut state = self.state.write(); state.cast_to_reader(vm)?; if let ReadState::CsvIter(ref mut reader) = &mut *state { if let Some(row) = reader.next() { match row { Ok(records) => { let iter = records .into_iter() .map(|bytes| bytes.into_pyobject(vm)) .collect::>(); Ok(vm.ctx.new_list(iter)) } Err(_err) => { let msg = String::from("Decode Error"); let decode_error = vm.new_unicode_decode_error(msg); Err(decode_error) } } } else { Err(vm.new_stop_iteration()) } } else { unreachable!() } } } fn _csv_reader(fp: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { if let Ok(iterable) = PyIterable::::try_from_object(vm, fp) { build_reader(iterable, args, vm) } else { Err(vm.new_type_error("argument 1 must be an iterator".to_owned())) } } pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; let reader_type = Reader::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, "Error" => error, // constants "QUOTE_MINIMAL" => ctx.new_int(QuoteStyle::QuoteMinimal as i32), "QUOTE_ALL" => ctx.new_int(QuoteStyle::QuoteAll as i32), "QUOTE_NONNUMERIC" => ctx.new_int(QuoteStyle::QuoteNonnumeric as i32), "QUOTE_NONE" => ctx.new_int(QuoteStyle::QuoteNone as i32), }) }