diff --git a/tests/snippets/stdlib_io_bytesio.py b/tests/snippets/stdlib_io_bytesio.py index f6f97e684..aa193dc9f 100644 --- a/tests/snippets/stdlib_io_bytesio.py +++ b/tests/snippets/stdlib_io_bytesio.py @@ -42,9 +42,21 @@ def test_04(): assert f.seek(0) == 0 assert f.read(4) == b'Test' +def test_05(): + """ + Tests that the write method accpets bytearray + """ + bytes_string = b'Test String 5' + + f = BytesIO() + f.write(bytearray(bytes_string)) + + assert f.getvalue() == bytes_string + if __name__ == "__main__": test_01() test_02() test_03() test_04() + test_05() diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index bee7e1d7f..fb59e0f06 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -11,6 +11,7 @@ use num_traits::ToPrimitive; use super::os; use crate::function::{OptionalArg, OptionalOption, PyFuncArgs}; use crate::obj::objbytearray::PyByteArray; +use crate::obj::objbyteinner::PyBytesLike; use crate::obj::objbytes; use crate::obj::objbytes::PyBytes; use crate::obj::objint::{self, PyIntRef}; @@ -173,11 +174,11 @@ impl PyValue for PyBytesIO { } impl PyBytesIORef { - fn write(self, data: objbytes::PyBytesRef, vm: &VirtualMachine) -> PyResult { - let bytes = data.get_value(); + fn write(self, data: PyBytesLike, vm: &VirtualMachine) -> PyResult { + let bytes = data.to_cow(); - match self.buffer.borrow_mut().write(bytes) { - Some(value) => Ok(vm.ctx.new_int(value)), + match self.buffer.borrow_mut().write(&bytes) { + Some(value) => Ok(value), None => Err(vm.new_type_error("Error Writing Bytes".to_string())), } }