From c0e1702869f21d3512307b4bb72195671ea73282 Mon Sep 17 00:00:00 2001 From: Maxim Kurnikov Date: Fri, 7 Jun 2019 16:34:16 +0300 Subject: [PATCH] split mode/type for io.open properly, fix open('r') case --- tests/snippets/builtin_open.py | 13 +- vm/src/stdlib/io.rs | 228 +++++++++++++++++++++++++-------- 2 files changed, 189 insertions(+), 52 deletions(-) diff --git a/tests/snippets/builtin_open.py b/tests/snippets/builtin_open.py index 41bcf2b80..d7c310a92 100644 --- a/tests/snippets/builtin_open.py +++ b/tests/snippets/builtin_open.py @@ -6,5 +6,14 @@ assert 'RustPython' in fd.read() assert_raises(FileNotFoundError, lambda: open('DoesNotExist')) # Use open as a context manager -with open('README.md') as fp: - fp.read() +with open('README.md', 'rt') as fp: + contents = fp.read() + assert type(contents) == str, "type is " + str(type(contents)) + +with open('README.md', 'r') as fp: + contents = fp.read() + assert type(contents) == str, "type is " + str(type(contents)) + +with open('README.md', 'rb') as fp: + contents = fp.read() + assert type(contents) == bytes, "type is " + str(type(contents)) diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index 8b203cc78..80b9ffb7f 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -1,9 +1,7 @@ /* * I/O core tools. */ - use std::cell::RefCell; -use std::collections::HashSet; use std::fs::File; use std::io::prelude::*; use std::io::BufReader; @@ -294,6 +292,63 @@ fn text_io_base_read(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { } } +fn split_mode_string(mode_string: String) -> Result<(String, String), String> { + let mut mode: char = '\0'; + let mut typ: char = '\0'; + let mut plus_is_set = false; + + for ch in mode_string.chars() { + match ch { + '+' => { + if plus_is_set { + return Err(format!("invalid mode: '{}'", mode_string)); + } + plus_is_set = true; + } + 't' | 'b' => { + if typ != '\0' { + if typ == ch { + // no duplicates allowed + return Err(format!("invalid mode: '{}'", mode_string)); + } else { + return Err("can't have text and binary mode at once".to_string()); + } + } + typ = ch; + } + 'a' | 'r' | 'w' => { + if mode != '\0' { + if mode == ch { + // no duplicates allowed + return Err(format!("invalid mode: '{}'", mode_string)); + } else { + return Err( + "must have exactly one of create/read/write/append mode".to_string() + ); + } + } + mode = ch; + } + _ => return Err(format!("invalid mode: '{}'", mode_string)), + } + } + + if mode == '\0' { + return Err( + "Must have exactly one of create/read/write/append mode and at most one plus" + .to_string(), + ); + } + let mut mode = mode.to_string(); + if plus_is_set { + mode.push('+'); + } + if typ == '\0' { + typ = 't'; + } + Ok((mode, typ.to_string())) +} + pub fn io_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, @@ -301,69 +356,55 @@ pub fn io_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { required = [(file, Some(vm.ctx.str_type()))], optional = [(mode, Some(vm.ctx.str_type()))] ); + // mode is optional: 'rt' is the default mode (open from reading text) + let mode_string = mode.map_or("rt".to_string(), |m| objstr::get_value(m)); + let (mode, typ) = match split_mode_string(mode_string) { + Ok((mode, typ)) => (mode, typ), + Err(error_message) => { + return Err(vm.new_value_error(error_message)); + } + }; let module = import::import_module(vm, PathBuf::default(), "io").unwrap(); - //mode is optional: 'rt' is the default mode (open from reading text) - let rust_mode = mode.map_or("rt".to_string(), |m| objstr::get_value(m)); - - let mut raw_modes = HashSet::new(); - - //add raw modes - raw_modes.insert("a".to_string()); - raw_modes.insert("r".to_string()); - raw_modes.insert("x".to_string()); - raw_modes.insert("w".to_string()); - - //This is not a terribly elegant way to separate the file mode from - //the "type" flag - this should be improved. The intention here is to - //match a valid flag for the file_io_init call: - //https://docs.python.org/3/library/io.html#io.FileIO - let modes: Vec = rust_mode - .chars() - .filter(|a| raw_modes.contains(&a.to_string())) - .collect(); - - if modes.is_empty() || modes.len() > 1 { - return Err(vm.new_value_error("Invalid Mode".to_string())); - } - - //Class objects (potentially) consumed by io.open - //RawIO: FileIO - //Buffered: BufferedWriter, BufferedReader - //Text: TextIOWrapper + // Class objects (potentially) consumed by io.open + // RawIO: FileIO + // Buffered: BufferedWriter, BufferedReader + // Text: TextIOWrapper let file_io_class = vm.get_attribute(module.clone(), "FileIO").unwrap(); let buffered_writer_class = vm.get_attribute(module.clone(), "BufferedWriter").unwrap(); let buffered_reader_class = vm.get_attribute(module.clone(), "BufferedReader").unwrap(); let text_io_wrapper_class = vm.get_attribute(module, "TextIOWrapper").unwrap(); - //Construct a FileIO (subclass of RawIOBase) - //This is subsequently consumed by a Buffered Class. - let file_args = vec![file.clone(), vm.ctx.new_str(modes[0].to_string())]; - let file_io = vm.invoke(file_io_class, file_args)?; + // Construct a FileIO (subclass of RawIOBase) + // This is subsequently consumed by a Buffered Class. + let file_args = vec![file.clone(), vm.ctx.new_str(mode.clone())]; + let file_io_obj = vm.invoke(file_io_class, file_args)?; - //Create Buffered class to consume FileIO. The type of buffered class depends on - //the operation in the mode. - //There are 3 possible classes here, each inheriting from the RawBaseIO + // Create Buffered class to consume FileIO. The type of buffered class depends on + // the operation in the mode. + // 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, vec![file_io.clone()]) - // reading => BufferedReader - } else { - vm.invoke(buffered_reader_class, vec![file_io.clone()]) + let buffered = match mode.chars().next().unwrap() { + 'w' => vm.invoke(buffered_writer_class, vec![file_io_obj.clone()]), + 'r' => vm.invoke(buffered_reader_class, vec![file_io_obj.clone()]), //TODO: updating => PyBufferedRandom + _ => unimplemented!("'a' mode is not yet implemented"), }; - 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, vec![buffered.unwrap()]) - } else { + let io_obj = match typ.chars().next().unwrap() { + // If the mode is text this buffer type is consumed on construction of + // a TextIOWrapper which is subsequently returned. + 't' => vm.invoke(text_io_wrapper_class, vec![buffered.unwrap()]), + // If the mode is binary this Buffered class is returned directly at // this point. - //For Buffered class construct "raw" IO class e.g. FileIO and pass this into corresponding field - buffered - } + // For Buffered class construct "raw" IO class e.g. FileIO and pass this into corresponding field + 'b' => buffered, + + _ => unreachable!(), + }; + io_obj } pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { @@ -439,3 +480,90 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "BytesIO" => bytes_io, }) } + +#[cfg(test)] +mod tests { + use super::*; + + fn assert_mode_split_into(mode_string: &str, expected_mode: &str, expected_typ: &str) { + let (mode, typ) = split_mode_string(mode_string.to_string()).unwrap(); + assert_eq!(mode, expected_mode); + assert_eq!(typ, expected_typ); + } + + #[test] + fn test_split_mode_valid_cases() { + assert_mode_split_into("r", "r", "t"); + assert_mode_split_into("rb", "r", "b"); + assert_mode_split_into("rt", "r", "t"); + assert_mode_split_into("r+t", "r+", "t"); + assert_mode_split_into("w+t", "w+", "t"); + assert_mode_split_into("r+b", "r+", "b"); + assert_mode_split_into("w+b", "w+", "b"); + } + + #[test] + fn test_invalid_mode() { + assert_eq!( + split_mode_string("rbsss".to_string()), + Err("invalid mode: 'rbsss'".to_string()) + ); + assert_eq!( + split_mode_string("rrb".to_string()), + Err("invalid mode: 'rrb'".to_string()) + ); + assert_eq!( + split_mode_string("rbb".to_string()), + Err("invalid mode: 'rbb'".to_string()) + ); + } + + #[test] + fn test_mode_not_specified() { + assert_eq!( + split_mode_string("".to_string()), + Err( + "Must have exactly one of create/read/write/append mode and at most one plus" + .to_string() + ) + ); + assert_eq!( + split_mode_string("b".to_string()), + Err( + "Must have exactly one of create/read/write/append mode and at most one plus" + .to_string() + ) + ); + assert_eq!( + split_mode_string("t".to_string()), + Err( + "Must have exactly one of create/read/write/append mode and at most one plus" + .to_string() + ) + ); + } + + #[test] + fn test_text_and_binary_at_once() { + assert_eq!( + split_mode_string("rbt".to_string()), + Err("can't have text and binary mode at once".to_string()) + ); + } + + #[test] + fn test_exactly_one_mode() { + assert_eq!( + split_mode_string("rwb".to_string()), + Err("must have exactly one of create/read/write/append mode".to_string()) + ); + } + + #[test] + fn test_at_most_one_plus() { + assert_eq!( + split_mode_string("a++".to_string()), + Err("invalid mode: 'a++'".to_string()) + ); + } +}