diff --git a/tests/snippets/test_io.py b/tests/snippets/test_io.py new file mode 100644 index 000000000..bc49b7287 --- /dev/null +++ b/tests/snippets/test_io.py @@ -0,0 +1,14 @@ + + +import io + +f = io.StringIO() +f.write('bladibla') +assert f.getvalue() == 'bladibla' + +# TODO: +# print('fubar', file=f, end='') +print(f.getvalue()) + +# TODO: +# assert f.getvalue() == 'bladiblafubar' diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index f9231ee23..aba31198a 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -2,6 +2,7 @@ * I/O core tools. */ +use std::cell::RefCell; use std::collections::HashSet; use std::fs::File; use std::io::prelude::*; @@ -19,7 +20,8 @@ use crate::obj::objbytes; use crate::obj::objint; use crate::obj::objstr; use crate::pyobject::{ - AttributeProtocol, BufferProtocol, PyContext, PyObjectRef, PyResult, TypeProtocol, + AttributeProtocol, BufferProtocol, PyContext, PyObject, PyObjectRef, PyRef, PyResult, PyValue, + TypeProtocol, }; use crate::vm::VirtualMachine; @@ -33,16 +35,39 @@ fn compute_c_flag(mode: &str) -> u16 { } } -fn string_io_init(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult { - // arg_check!(vm, args, required = [(s, Some(vm.ctx.str_type()))]); - // TODO - Ok(vm.get_none()) +#[derive(Debug)] +struct PyStringIO { + data: RefCell, } -fn string_io_getvalue(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { - arg_check!(vm, args); - // TODO - Ok(vm.get_none()) +type PyStringIORef = PyRef; + +impl PyValue for PyStringIO { + fn class(vm: &mut VirtualMachine) -> PyObjectRef { + vm.class("io", "StringIO") + } +} + +impl PyStringIORef { + fn write(self, data: objstr::PyStringRef, _vm: &mut VirtualMachine) { + let data = data.value.clone(); + self.data.borrow_mut().push_str(&data); + } + + fn getvalue(self, _vm: &mut VirtualMachine) -> String { + self.data.borrow().clone() + } +} + +fn string_io_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { + arg_check!(vm, args, required = [(cls, None)]); + + Ok(PyObject::new( + PyStringIO { + data: RefCell::new(String::default()), + }, + cls.clone(), + )) } fn bytes_io_init(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult { @@ -97,8 +122,7 @@ fn buffered_reader_read(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { //to obtain buff_size many bytes. Exit when less than buff_size many //bytes are returned (when the end of the file is reached). while length == buff_size { - let raw_read = vm.get_method(raw.clone(), &"readinto".to_string()).unwrap(); - vm.invoke(raw_read, PyFuncArgs::new(vec![buffer.clone()], vec![])) + vm.call_method(&raw, "readinto", vec![buffer.clone()]) .map_err(|_| vm.new_value_error("IO Error".to_string()))?; //Copy bytes from the buffer vector into the results vector @@ -106,9 +130,8 @@ fn buffered_reader_read(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { result.extend_from_slice(&bytes.value.borrow()); }; - let len = vm.get_method(buffer.clone(), &"__len__".to_string()); - let py_len = vm.invoke(len.unwrap(), PyFuncArgs::default()); - length = objint::get_value(&py_len.unwrap()).to_usize().unwrap(); + let py_len = vm.call_method(&buffer, "__len__", PyFuncArgs::default())?; + length = objint::get_value(&py_len).to_usize().unwrap(); } Ok(vm.ctx.new_bytes(result)) @@ -174,9 +197,8 @@ fn file_io_readinto(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { } //extract length of buffer - let len_method = vm.get_method(obj.clone(), &"__len__".to_string()); - let py_length = vm.invoke(len_method.unwrap(), PyFuncArgs::default()); - let length = objint::get_value(&py_length.unwrap()).to_u64().unwrap(); + let py_length = vm.call_method(obj, "__len__", PyFuncArgs::default())?; + let length = objint::get_value(&py_length).to_u64().unwrap(); let file_no = file_io.get_attr("fileno").unwrap(); let raw_fd = objint::get_value(&file_no).to_i64().unwrap(); @@ -243,10 +265,9 @@ fn buffered_writer_write(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult ); let raw = vm.ctx.get_attr(&buffered, "raw").unwrap(); - let raw_write = vm.get_method(raw.clone(), &"write".to_string()).unwrap(); //This should be replaced with a more appropriate chunking implementation - vm.invoke(raw_write, PyFuncArgs::new(vec![obj.clone()], vec![])) + vm.call_method(&raw, "write", vec![obj.clone()]) } fn text_io_wrapper_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { @@ -264,9 +285,8 @@ fn text_io_base_read(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(text_io_base, None)]); let raw = vm.ctx.get_attr(&text_io_base, "buffer").unwrap(); - let read = vm.get_method(raw.clone(), &"read".to_string()); - if let Ok(bytes) = vm.invoke(read.unwrap(), PyFuncArgs::default()) { + if let Ok(bytes) = vm.call_method(&raw, "read", PyFuncArgs::default()) { let value = objbytes::get_value(&bytes).to_vec(); //format bytes into string @@ -322,10 +342,7 @@ pub fn io_open(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { //Construct a FileIO (subclass of RawIOBase) //This is subsequently consumed by a Buffered Class. - let file_args = PyFuncArgs::new( - vec![file.clone(), vm.ctx.new_str(modes[0].to_string())], - vec![], - ); + let file_args = vec![file.clone(), vm.ctx.new_str(modes[0].to_string())]; let file_io = vm.invoke(file_io_class, file_args)?; //Create Buffered class to consume FileIO. The type of buffered class depends on @@ -333,26 +350,17 @@ pub fn io_open(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult { //There are 3 possible classes here, each inheriting from the RawBaseIO // creating || writing || appending => BufferedWriter let buffered = if rust_mode.contains('w') { - vm.invoke( - buffered_writer_class, - PyFuncArgs::new(vec![file_io.clone()], vec![]), - ) + vm.invoke(buffered_writer_class, vec![file_io.clone()]) // reading => BufferedReader } else { - vm.invoke( - buffered_reader_class, - PyFuncArgs::new(vec![file_io.clone()], vec![]), - ) + vm.invoke(buffered_reader_class, vec![file_io.clone()]) //TODO: updating => PyBufferedRandom }; if rust_mode.contains('t') { //If the mode is text this buffer type is consumed on construction of //a TextIOWrapper which is subsequently returned. - vm.invoke( - text_io_wrapper_class, - PyFuncArgs::new(vec![buffered.unwrap()], vec![]), - ) + vm.invoke(text_io_wrapper_class, vec![buffered.unwrap()]) } else { // If the mode is binary this Buffered class is returned directly at // this point. @@ -405,8 +413,9 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef { //StringIO: in-memory text let string_io = py_class!(ctx, "StringIO", text_io_base.clone(), { - "__init__" => ctx.new_rustfunc(string_io_init), - "getvalue" => ctx.new_rustfunc(string_io_getvalue) + "__new__" => ctx.new_rustfunc(string_io_new), + "write" => ctx.new_rustfunc(PyStringIORef::write), + "getvalue" => ctx.new_rustfunc(PyStringIORef::getvalue) }); //BytesIO: in-memory bytes