Merge pull request #673 from RustPython/string-io

Add write method to io.StringIO

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'

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<String>,
 }
 
-fn string_io_getvalue(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    arg_check!(vm, args);
-    // TODO
-    Ok(vm.get_none())
+type PyStringIORef = PyRef<PyStringIO>;
+
+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