/* * I/O core tools. */ use std::cell::RefCell; use std::collections::HashSet; use std::fs::File; use std::io::prelude::*; use std::io::BufReader; use std::path::PathBuf; use num_bigint::ToBigInt; use num_traits::ToPrimitive; use super::os; use crate::function::PyFuncArgs; use crate::import; use crate::obj::objbytearray::PyByteArray; use crate::obj::objbytes; use crate::obj::objint; use crate::obj::objstr; use crate::obj::objtype::PyClassRef; use crate::pyobject::{ BufferProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, }; use crate::vm::VirtualMachine; fn compute_c_flag(mode: &str) -> u16 { match mode { "w" => 512, "x" => 512, "a" => 8, "+" => 2, _ => 0, } } #[derive(Debug)] struct PyStringIO { data: RefCell, } type PyStringIORef = PyRef; impl PyValue for PyStringIO { fn class(vm: &VirtualMachine) -> PyClassRef { vm.class("io", "StringIO") } } impl PyStringIORef { fn write(self, data: objstr::PyStringRef, _vm: &VirtualMachine) { let data = data.value.clone(); self.data.borrow_mut().push_str(&data); } fn getvalue(self, _vm: &VirtualMachine) -> String { self.data.borrow().clone() } } fn string_io_new(cls: PyClassRef, vm: &VirtualMachine) -> PyResult { PyStringIO { data: RefCell::new(String::default()), } .into_ref_with_type(vm, cls) } fn bytes_io_init(vm: &VirtualMachine, _args: PyFuncArgs) -> PyResult { // TODO Ok(vm.get_none()) } fn bytes_io_getvalue(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args); // TODO Ok(vm.get_none()) } fn io_base_cm_enter(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(instance, None)]); Ok(instance.clone()) } fn io_base_cm_exit(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, args, // The context manager protocol requires these, but we don't use them required = [ (_instance, None), (_exception_type, None), (_exception_value, None), (_traceback, None) ] ); Ok(vm.get_none()) } fn buffered_io_base_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(buffered, None), (raw, None)]); vm.set_attr(buffered, "raw", raw.clone())?; Ok(vm.get_none()) } fn buffered_reader_read(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(buffered, None)]); let buff_size = 8 * 1024; let buffer = vm.ctx.new_bytearray(vec![0; buff_size]); //buffer method let mut result = vec![]; let mut length = buff_size; let raw = vm.get_attribute(buffered.clone(), "raw").unwrap(); //Iterates through the raw class, invoking the readinto method //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 { 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 if let Some(bytes) = buffer.payload::() { result.extend_from_slice(&bytes.value.borrow()); }; 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)) } fn file_io_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, args, required = [(file_io, None), (name, Some(vm.ctx.str_type()))], optional = [(mode, Some(vm.ctx.str_type()))] ); let rust_mode = mode.map_or("r".to_string(), |m| objstr::get_value(m)); match compute_c_flag(&rust_mode).to_bigint() { Some(os_mode) => { let args = vec![name.clone(), vm.ctx.new_int(os_mode)]; let file_no = os::os_open(vm, PyFuncArgs::new(args, vec![]))?; vm.set_attr(file_io, "name", name.clone())?; vm.set_attr(file_io, "fileno", file_no)?; vm.set_attr(file_io, "closefd", vm.new_bool(false))?; vm.set_attr(file_io, "closed", vm.new_bool(false))?; Ok(vm.get_none()) } None => Err(vm.new_type_error(format!("invalid mode {}", rust_mode))), } } fn file_io_read(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(file_io, None)]); let py_name = vm.get_attribute(file_io.clone(), "name")?; let f = match File::open(objstr::get_value(&py_name)) { Ok(v) => Ok(v), Err(_) => Err(vm.new_type_error("Error opening file".to_string())), }; let buffer = match f { Ok(v) => Ok(BufReader::new(v)), Err(_) => Err(vm.new_type_error("Error reading from file".to_string())), }; let mut bytes = vec![]; if let Ok(mut buff) = buffer { match buff.read_to_end(&mut bytes) { Ok(_) => {} Err(_) => return Err(vm.new_value_error("Error reading from Buffer".to_string())), } } Ok(vm.ctx.new_bytes(bytes)) } fn file_io_readinto(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(file_io, None), (obj, None)]); if !obj.readonly() { return Ok(vm.new_type_error( "readinto() argument must be read-write bytes-like object".to_string(), )); } //extract length of buffer let py_length = vm.call_method(obj, "__len__", PyFuncArgs::default())?; let length = objint::get_value(&py_length).to_u64().unwrap(); let file_no = vm.get_attribute(file_io.clone(), "fileno")?; let raw_fd = objint::get_value(&file_no).to_i64().unwrap(); //extract unix file descriptor. let handle = os::rust_file(raw_fd); let mut f = handle.take(length); if let Some(bytes) = obj.payload::() { //TODO: Implement for MemoryView let mut value_mut = bytes.value.borrow_mut(); value_mut.clear(); match f.read_to_end(&mut value_mut) { Ok(_) => {} Err(_) => return Err(vm.new_value_error("Error reading from Take".to_string())), } }; let updated = os::raw_file_number(f.into_inner()); vm.set_attr(file_io, "fileno", vm.ctx.new_int(updated))?; Ok(vm.get_none()) } fn file_io_write(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, args, required = [(file_io, None), (obj, Some(vm.ctx.bytes_type()))] ); let file_no = vm.get_attribute(file_io.clone(), "fileno")?; let raw_fd = objint::get_value(&file_no).to_i64().unwrap(); //unsafe block - creates file handle from the UNIX file descriptor //raw_fd is supported on UNIX only. This will need to be extended //to support windows - i.e. raw file_handles let mut handle = os::rust_file(raw_fd); match obj.payload::() { Some(bytes) => { let value_mut = bytes.value.borrow(); match handle.write(&value_mut[..]) { Ok(len) => { //reset raw fd on the FileIO object let updated = os::raw_file_number(handle); vm.set_attr(file_io, "fileno", vm.ctx.new_int(updated))?; //return number of bytes written Ok(vm.ctx.new_int(len)) } Err(_) => Err(vm.new_value_error("Error Writing Bytes to Handle".to_string())), } } None => Err(vm.new_value_error("Expected Bytes Object".to_string())), } } fn buffered_writer_write(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, args, required = [(buffered, None), (obj, Some(vm.ctx.bytes_type()))] ); let raw = vm.get_attribute(buffered.clone(), "raw").unwrap(); //This should be replaced with a more appropriate chunking implementation vm.call_method(&raw, "write", vec![obj.clone()]) } fn text_io_wrapper_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, args, required = [(text_io_wrapper, None), (buffer, None)] ); vm.set_attr(text_io_wrapper, "buffer", buffer.clone())?; Ok(vm.get_none()) } fn text_io_base_read(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!(vm, args, required = [(text_io_base, None)]); let raw = vm.get_attribute(text_io_base.clone(), "buffer").unwrap(); 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(); Ok(vm.ctx.new_str(rust_string)) } else { Err(vm.new_value_error("Error unpacking Bytes".to_string())) } } pub fn io_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { arg_check!( vm, args, required = [(file, Some(vm.ctx.str_type()))], optional = [(mode, Some(vm.ctx.str_type()))] ); 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 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)?; //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()]) //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, vec![buffered.unwrap()]) } else { // 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 } } pub fn make_module(ctx: &PyContext) -> PyObjectRef { //IOBase the abstract base class of the IO Module let io_base = py_class!(ctx, "IOBase", ctx.object(), { "__enter__" => ctx.new_rustfunc(io_base_cm_enter), "__exit__" => ctx.new_rustfunc(io_base_cm_exit) }); // IOBase Subclasses let raw_io_base = py_class!(ctx, "RawIOBase", ctx.object(), {}); let buffered_io_base = py_class!(ctx, "BufferedIOBase", io_base.clone(), { "__init__" => ctx.new_rustfunc(buffered_io_base_init) }); //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) }); // RawBaseIO Subclasses // TODO Fix name? let file_io = py_class!(ctx, "FileIO", raw_io_base.clone(), { "__init__" => ctx.new_rustfunc(file_io_init), "name" => ctx.str_type(), "read" => ctx.new_rustfunc(file_io_read), "readinto" => ctx.new_rustfunc(file_io_readinto), "write" => ctx.new_rustfunc(file_io_write) }); // BufferedIOBase Subclasses let buffered_reader = py_class!(ctx, "BufferedReader", buffered_io_base.clone(), { "read" => ctx.new_rustfunc(buffered_reader_read) }); let buffered_writer = py_class!(ctx, "BufferedWriter", buffered_io_base.clone(), { "write" => ctx.new_rustfunc(buffered_writer_write) }); //TextIOBase Subclass let text_io_wrapper = py_class!(ctx, "TextIOWrapper", text_io_base.clone(), { "__init__" => ctx.new_rustfunc(text_io_wrapper_init) }); //StringIO: in-memory text let string_io = py_class!(ctx, "StringIO", text_io_base.clone(), { "__new__" => ctx.new_rustfunc(string_io_new), "write" => ctx.new_rustfunc(PyStringIORef::write), "getvalue" => ctx.new_rustfunc(PyStringIORef::getvalue) }); //BytesIO: in-memory bytes let bytes_io = py_class!(ctx, "BytesIO", buffered_io_base.clone(), { "__init__" => ctx.new_rustfunc(bytes_io_init), "getvalue" => ctx.new_rustfunc(bytes_io_getvalue) }); py_module!(ctx, "io", { "open" => ctx.new_rustfunc(io_open), "IOBase" => io_base, "RawIOBase" => raw_io_base, "BufferedIOBase" => buffered_io_base, "TextIOBase" => text_io_base, "FileIO" => file_io, "BufferedReader" => buffered_reader, "BufferedWriter" => buffered_writer, "TextIOWrapper" => text_io_wrapper, "StringIO" => string_io, "BytesIO" => bytes_io, }) }