diff --git a/tests/snippets/stdlib_os.py b/tests/snippets/stdlib_os.py index ed9e91c70..ba6033ff8 100644 --- a/tests/snippets/stdlib_os.py +++ b/tests/snippets/stdlib_os.py @@ -2,7 +2,36 @@ import os from testutils import assert_raises -assert os.open('README.md', 0) > 0 +fd = os.open('README.md', 0) +assert fd > 0 + +os.close(fd) +assert_raises(OSError, lambda: os.read(fd, 10)) + +FNAME = "test_file_that_no_one_will_have_on_disk" +CONTENT = b"testing" +CONTENT2 = b"rustpython" +CONTENT3 = b"BOYA" + +class TestWithFile(): + def __enter__(self): + open(FNAME, "wb") + return FNAME + + def __exit__(self, exc_type, exc_val, exc_tb): + os.remove(FNAME) + + +with TestWithFile() as fname: + fd = os.open(fname, 1) + assert os.write(fd, CONTENT2) == len(CONTENT2) + assert os.write(fd, CONTENT3) == len(CONTENT3) + os.close(fd) + + fd = os.open(fname, 0) + assert os.read(fd, len(CONTENT2)) == CONTENT2 + assert os.read(fd, len(CONTENT3)) == CONTENT3 + os.close(fd) assert_raises(FileNotFoundError, lambda: os.open('DOES_NOT_EXIST', 0)) diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs index f24c9273c..4a8b1ad5e 100644 --- a/vm/src/obj/objbytes.rs +++ b/vm/src/obj/objbytes.rs @@ -17,7 +17,7 @@ use super::objtype::PyClassRef; pub struct PyBytes { value: Vec, } -type PyBytesRef = PyRef; +pub type PyBytesRef = PyRef; impl PyBytes { pub fn new(data: Vec) -> Self { diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index 3f0f557f1..ae487598a 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -1,12 +1,16 @@ +use std::fs; use std::fs::File; use std::fs::OpenOptions; -use std::io::ErrorKind; +use std::io::{ErrorKind, Read, Write}; use num_traits::cast::ToPrimitive; use crate::function::PyFuncArgs; +use crate::obj::objbytes::PyBytesRef; use crate::obj::objint; +use crate::obj::objint::PyIntRef; use crate::obj::objstr; +use crate::obj::objstr::PyStringRef; use crate::pyobject::{PyObjectRef, PyResult, TypeProtocol}; use crate::vm::VirtualMachine; @@ -113,6 +117,40 @@ fn os_error(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { Err(vm.new_os_error(msg)) } +fn os_read(fd: PyIntRef, n: PyIntRef, vm: &VirtualMachine) -> PyResult { + let mut buffer = vec![0u8; n.as_bigint().to_usize().unwrap()]; + let mut file = rust_file(fd.as_bigint().to_i64().unwrap()); + match file.read_exact(&mut buffer) { + Ok(_) => (), + Err(s) => return Err(vm.new_os_error(s.to_string())), + }; + + // Avoid closing the fd + raw_file_number(file); + Ok(vm.ctx.new_bytes(buffer)) +} + +fn os_write(fd: PyIntRef, data: PyBytesRef, vm: &VirtualMachine) -> PyResult { + let mut file = rust_file(fd.as_bigint().to_i64().unwrap()); + let written = match file.write(&data) { + Ok(written) => written, + Err(s) => return Err(vm.new_os_error(s.to_string())), + }; + + // Avoid closing the fd + raw_file_number(file); + Ok(vm.ctx.new_int(written)) +} + +fn os_remove(path: PyStringRef, vm: &VirtualMachine) -> PyResult { + match fs::remove_file(&path.value) { + Ok(_) => (), + Err(s) => return Err(vm.new_os_error(s.to_string())), + } + + Ok(vm.get_none()) +} + pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; @@ -126,6 +164,10 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "open" => ctx.new_rustfunc(os_open), "close" => ctx.new_rustfunc(os_close), "error" => ctx.new_rustfunc(os_error), + "read" => ctx.new_rustfunc(os_read), + "write" => ctx.new_rustfunc(os_write), + "remove" => ctx.new_rustfunc(os_remove), + "unlink" => ctx.new_rustfunc(os_remove), "name" => ctx.new_str(os_name), "O_RDONLY" => ctx.new_int(0), "O_WRONLY" => ctx.new_int(1),