Merge pull request RustPython#1119 from silmeth/text_file_io

make TextIOBase writable, handle malformed utf-8 in read()

Author: windelbouwman

vm/src/builtins.rs

@@ -889,6 +889,10 @@ pub fn make_module(vm: &VirtualMachine, module: PyObjectRef) {
         "FileNotFoundError" => ctx.exceptions.file_not_found_error.clone(),
         "FileExistsError" => ctx.exceptions.file_exists_error.clone(),
         "StopIteration" => ctx.exceptions.stop_iteration.clone(),
+        "UnicodeError" => ctx.exceptions.unicode_error.clone(),
+        "UnicodeDecodeError" => ctx.exceptions.unicode_decode_error.clone(),
+        "UnicodeEncodeError" => ctx.exceptions.unicode_encode_error.clone(),
+        "UnicodeTranslateError" => ctx.exceptions.unicode_translate_error.clone(),
         "ZeroDivisionError" => ctx.exceptions.zero_division_error.clone(),
         "KeyError" => ctx.exceptions.key_error.clone(),
         "OSError" => ctx.exceptions.os_error.clone(),

vm/src/exceptions.rs

@@ -215,6 +215,10 @@ pub struct ExceptionZoo {
     pub syntax_error: PyClassRef,
     pub type_error: PyClassRef,
     pub value_error: PyClassRef,
+    pub unicode_error: PyClassRef,
+    pub unicode_decode_error: PyClassRef,
+    pub unicode_encode_error: PyClassRef,
+    pub unicode_translate_error: PyClassRef,
     pub zero_division_error: PyClassRef,
     pub eof_error: PyClassRef,
 
@@ -258,6 +262,11 @@ impl ExceptionZoo {
         let permission_error = create_type("PermissionError", &type_type, &os_error);
         let file_exists_error = create_type("FileExistsError", &type_type, &os_error);
         let eof_error = create_type("EOFError", &type_type, &exception_type);
+        let unicode_error = create_type("UnicodeError", &type_type, &value_error);
+        let unicode_decode_error = create_type("UnicodeDecodeError", &type_type, &unicode_error);
+        let unicode_encode_error = create_type("UnicodeEncodeError", &type_type, &unicode_error);
+        let unicode_translate_error =
+            create_type("UnicodeTranslateError", &type_type, &unicode_error);
 
         let warning = create_type("Warning", &type_type, &exception_type);
         let bytes_warning = create_type("BytesWarning", &type_type, &warning);
@@ -294,6 +303,10 @@ impl ExceptionZoo {
             syntax_error,
             type_error,
             value_error,
+            unicode_error,
+            unicode_decode_error,
+            unicode_encode_error,
+            unicode_translate_error,
             zero_division_error,
             eof_error,
             warning,

vm/src/stdlib/io.rs

@@ -442,19 +442,62 @@ fn text_io_wrapper_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
 fn text_io_base_read(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(text_io_base, None)]);
 
+    let buffered_reader_class = vm.try_class("_io", "BufferedReader")?;
     let raw = vm.get_attribute(text_io_base.clone(), "buffer").unwrap();
 
+    if !objtype::isinstance(&raw, &buffered_reader_class) {
+        // TODO: this should be io.UnsupportedOperation error which derives both from ValueError *and* OSError
+        return Err(vm.new_value_error("not readable".to_string()));
+    }
+
     if let Ok(bytes) = vm.call_method(&raw, "read", PyFuncArgs::default()) {
         let value = objbytes::get_value(&bytes).to_vec();
 
         //format bytes into string
-        let rust_string = String::from_utf8(value).unwrap();
+        let rust_string = String::from_utf8(value).map_err(|e| {
+            vm.new_unicode_decode_error(format!(
+                "cannot decode byte at index: {}",
+                e.utf8_error().valid_up_to()
+            ))
+        })?;
         Ok(vm.ctx.new_str(rust_string))
     } else {
         Err(vm.new_value_error("Error unpacking Bytes".to_string()))
     }
 }
 
+fn text_io_base_write(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    use std::str::from_utf8;
+
+    arg_check!(
+        vm,
+        args,
+        required = [(text_io_base, None), (obj, Some(vm.ctx.str_type()))]
+    );
+
+    let buffered_writer_class = vm.try_class("_io", "BufferedWriter")?;
+    let raw = vm.get_attribute(text_io_base.clone(), "buffer").unwrap();
+
+    if !objtype::isinstance(&raw, &buffered_writer_class) {
+        // TODO: this should be io.UnsupportedOperation error which derives from ValueError and OSError
+        return Err(vm.new_value_error("not writable".to_string()));
+    }
+
+    let bytes = objstr::get_value(obj).into_bytes();
+
+    let len = vm.call_method(&raw, "write", vec![vm.ctx.new_bytes(bytes.clone())])?;
+    let len = objint::get_value(&len).to_usize().ok_or_else(|| {
+        vm.new_overflow_error("int to large to convert to Rust usize".to_string())
+    })?;
+
+    // returns the count of unicode code points written
+    let len = from_utf8(&bytes[..len])
+        .unwrap_or_else(|e| from_utf8(&bytes[..e.valid_up_to()]).unwrap())
+        .chars()
+        .count();
+    Ok(vm.ctx.new_int(len))
+}
+
 fn split_mode_string(mode_string: String) -> Result<(String, String), String> {
     let mut mode: char = '\0';
     let mut typ: char = '\0';
@@ -594,7 +637,8 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
 
     //TextIO Base has no public constructor
     let text_io_base = py_class!(ctx, "TextIOBase", io_base.clone(), {
-        "read" => ctx.new_rustfunc(text_io_base_read)
+        "read" => ctx.new_rustfunc(text_io_base_read),
+        "write" => ctx.new_rustfunc(text_io_base_write)
     });
 
     // RawBaseIO Subclasses

vm/src/vm.rs

@@ -219,6 +219,11 @@ impl VirtualMachine {
         self.new_exception(os_error, msg)
     }
 
+    pub fn new_unicode_decode_error(&self, msg: String) -> PyObjectRef {
+        let unicode_decode_error = self.ctx.exceptions.unicode_decode_error.clone();
+        self.new_exception(unicode_decode_error, msg)
+    }
+
     /// Create a new python ValueError object. Useful for raising errors from
     /// python functions implemented in rust.
     pub fn new_value_error(&self, msg: String) -> PyObjectRef {