diff --git a/crates/vm/src/stdlib/_io.rs b/crates/vm/src/stdlib/_io.rs index 6c1fa064a..69f57e32a 100644 --- a/crates/vm/src/stdlib/_io.rs +++ b/crates/vm/src/stdlib/_io.rs @@ -4772,12 +4772,85 @@ mod _io { #[derive(Debug)] struct StringIO { _base: _TextIOBase, - buffer: PyRwLock, + buffer: PyRwLock, newline: AtomicCell, seennl: AtomicCell, closed: AtomicCell, } + #[derive(Debug)] + struct StringIOBuffer { + bytes: BufferedIO, + position: usize, + length: usize, + } + + impl StringIOBuffer { + fn new(bytes: Vec) -> Self { + let length = StringIO::text(&bytes).code_points().count(); + Self { + bytes: BufferedIO::new(Cursor::new(bytes)), + position: 0, + length, + } + } + + fn byte_offset(&self, position: usize) -> usize { + let bytes = self.bytes.cursor.get_ref(); + if position >= self.length { + bytes.len() + (position - self.length) + } else { + crate::common::str::codepoint_range_end(StringIO::text(bytes), position).unwrap() + } + } + + fn seek(&mut self, position: usize) { + let byte_position = self.byte_offset(position); + self.bytes.cursor.set_position(byte_position as u64); + self.position = position; + } + + fn read(&mut self, size: usize) -> Vec { + let data = self.bytes.read(Some(size)).unwrap_or_default(); + self.position += StringIO::text(&data).code_points().count(); + data + } + + fn write(&mut self, data: &Wtf8, count: usize, vm: &VirtualMachine) -> PyResult<()> { + if data.is_empty() { + return Ok(()); + } + let position = self + .position + .checked_add(count) + .filter(|&position| isize::try_from(position).is_ok()) + .ok_or_else(|| vm.new_overflow_error("new position too large"))?; + let byte_position = self.bytes.tell() as usize; + let bytes = self.bytes.cursor.get_mut(); + if byte_position > bytes.len() { + bytes.resize(byte_position, 0); + } + let tail = StringIO::text(&bytes[byte_position..]); + let replaced = if count >= self.length.saturating_sub(self.position) { + tail.len() + } else if tail.as_bytes()[..count].is_ascii() { + count + } else { + crate::common::str::codepoint_range_end(tail, count).unwrap() + }; + bytes.splice( + byte_position..byte_position + replaced, + data.as_bytes().iter().copied(), + ); + self.bytes + .cursor + .set_position((byte_position + data.len()) as u64); + self.position = position; + self.length = self.length.max(position); + Ok(()) + } + } + #[derive(FromArgs)] struct StringIONewArgs { // None means an empty buffer. @@ -4795,7 +4868,7 @@ mod _io { fn py_new(_cls: &Py, _args: Self::Args, _vm: &VirtualMachine) -> PyResult { Ok(Self { _base: Default::default(), - buffer: PyRwLock::new(BufferedIO::new(Cursor::new(Vec::new()))), + buffer: PyRwLock::new(StringIOBuffer::new(Vec::new())), newline: AtomicCell::new(Newlines::Lf), seennl: AtomicCell::new(SeenNewline::empty()), closed: AtomicCell::new(false), @@ -4822,7 +4895,7 @@ mod _io { let raw_bytes = initial_value.as_ref().map_or_else(Vec::new, |v| { Self::translate_newlines(v.as_wtf8(), newline).into_bytes() }); - *zelf.buffer.write() = BufferedIO::new(Cursor::new(raw_bytes)); + *zelf.buffer.write() = StringIOBuffer::new(raw_bytes); zelf.newline.store(newline); zelf.seennl.store(SeenNewline::empty()); if let Some(initial_value) = initial_value { @@ -4833,7 +4906,7 @@ mod _io { } impl StringIO { - fn buffer(&self, vm: &VirtualMachine) -> PyResult> { + fn buffer(&self, vm: &VirtualMachine) -> PyResult> { if !self.closed.load() { Ok(self.buffer.write()) } else { @@ -4865,19 +4938,6 @@ mod _io { unsafe { Wtf8::from_bytes_unchecked(bytes) } } - fn char_offset_to_byte(bytes: &[u8], char_offset: usize) -> usize { - let text = Self::text(bytes); - crate::common::str::codepoint_range_end(text, char_offset) - .unwrap_or_else(|| bytes.len() + (char_offset - text.code_points().count())) - } - - fn byte_offset_to_char(bytes: &[u8], byte_offset: usize) -> usize { - let content_len = bytes.len(); - let in_content = byte_offset.min(content_len); - Self::text(&bytes[..in_content]).code_points().count() - + byte_offset.saturating_sub(content_len) - } - fn read_size(buffer: &BufferedIO, size: Option, newline: Option) -> usize { let position = buffer.tell() as usize; let bytes = buffer.cursor.get_ref().get(position..).unwrap_or_default(); @@ -4926,23 +4986,26 @@ mod _io { self.closed.store(true); } - // write string to underlying vector + // Replace characters in the underlying vector, regardless of their byte widths. #[pymethod] fn write(&self, s: PyStrRef, vm: &VirtualMachine) -> PyResult { let newline = self.newline.load(); - let bytes = StringIO::translate_newlines(s.as_wtf8(), newline).into_bytes(); + let text = StringIO::translate_newlines(s.as_wtf8(), newline); + let count = if matches!(newline, Newlines::Universal | Newlines::Crlf) { + text.code_points().count() + } else { + s.char_len() + }; let mut buffer = self.buffer(vm)?; self.observe_newlines(s.as_wtf8(), newline); - buffer - .write(&bytes) - .ok_or_else(|| vm.new_type_error("Error Writing String"))?; + buffer.write(&text, count, vm)?; Ok(s.char_len() as u64) } // return the entire contents of the underlying #[pymethod] fn getvalue(&self, vm: &VirtualMachine) -> PyResult { - let bytes = self.buffer(vm)?.getvalue(); + let bytes = self.buffer(vm)?.bytes.getvalue(); Wtf8Buf::from_bytes(bytes).map_err(|_| vm.new_value_error("Error Retrieving Value")) } @@ -4959,14 +5022,8 @@ mod _io { let kind = if how == 1 { "cur" } else { "end" }; return Err(vm.new_os_error(format!("can't do nonzero {kind}-relative seeks"))); } - 1 | 2 => { - let byte_offset = if how == 1 { - buffer.tell() as usize - } else { - buffer.cursor.get_ref().len() - }; - StringIO::byte_offset_to_char(buffer.cursor.get_ref(), byte_offset) - } + 1 => return Ok(buffer.position as u64), + 2 => buffer.length, _ => { return Err( vm.new_value_error(format!("invalid whence ({how}, should be 0, 1 or 2)")) @@ -4974,10 +5031,7 @@ mod _io { } }; - let byte_offset = StringIO::char_offset_to_byte(buffer.cursor.get_ref(), char_offset); - buffer - .seek(SeekFrom::Start(byte_offset as u64)) - .map_err(|err| os_err(vm, err))?; + buffer.seek(char_offset); Ok(char_offset as u64) } @@ -4987,8 +5041,8 @@ mod _io { #[pymethod] fn read(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult { let mut buffer = self.buffer(vm)?; - let size = StringIO::read_size(&buffer, size.to_usize(), None); - let data = buffer.read(Some(size)).unwrap_or_default(); + let size = StringIO::read_size(&buffer.bytes, size.to_usize(), None); + let data = buffer.read(size); let value = Wtf8Buf::from_bytes(data) .map_err(|_| vm.new_value_error("Error Retrieving Value"))?; @@ -4997,32 +5051,30 @@ mod _io { #[pymethod] fn tell(&self, vm: &VirtualMachine) -> PyResult { - let buffer = self.buffer(vm)?; - Ok( - StringIO::byte_offset_to_char(buffer.cursor.get_ref(), buffer.tell() as usize) - as u64, - ) + Ok(self.buffer(vm)?.position as u64) } #[pymethod] fn readline(&self, size: OptionalSize, vm: &VirtualMachine) -> PyResult { let mut buffer = self.buffer(vm)?; - let size = StringIO::read_size(&buffer, size.to_usize(), Some(self.newline.load())); - let input = buffer.read(Some(size)).unwrap_or_default(); + let size = + StringIO::read_size(&buffer.bytes, size.to_usize(), Some(self.newline.load())); + let input = buffer.read(size); Wtf8Buf::from_bytes(input).map_err(|_| vm.new_value_error("Error Retrieving Value")) } #[pymethod] fn truncate(&self, pos: StringIOPos, vm: &VirtualMachine) -> PyResult { let mut buffer = self.buffer(vm)?; - let pos = match pos.as_optional().try_usize(vm)? { - Some(pos) => pos, - None => { - StringIO::byte_offset_to_char(buffer.cursor.get_ref(), buffer.tell() as usize) - } - }; - let byte_pos = StringIO::char_offset_to_byte(buffer.cursor.get_ref(), pos); - buffer.truncate(Some(byte_pos)); + let pos = pos.as_optional().try_usize(vm)?.unwrap_or(buffer.position); + let byte_pos = buffer.byte_offset(pos); + buffer.bytes.truncate(Some(byte_pos)); + buffer.length = buffer.length.min(pos); + // Truncation keeps the character position, including positions beyond EOF. + if buffer.bytes.tell() as usize > buffer.bytes.cursor.get_ref().len() { + let position = buffer.position; + buffer.seek(position); + } Ok(pos) } @@ -5034,10 +5086,9 @@ mod _io { #[pymethod] fn __getstate__(zelf: PyRef, vm: &VirtualMachine) -> PyResult { let buffer = zelf.buffer(vm)?; - let content = Wtf8Buf::from_bytes(buffer.getvalue()) + let content = Wtf8Buf::from_bytes(buffer.bytes.getvalue()) .map_err(|_| vm.new_value_error("Error Retrieving Value"))?; - let pos = - StringIO::byte_offset_to_char(buffer.cursor.get_ref(), buffer.tell() as usize); + let pos = buffer.position; drop(buffer); // Get __dict__ if it exists and is non-empty @@ -5091,11 +5142,8 @@ mod _io { // Set content and position let raw_bytes = content.as_bytes().to_vec(); let mut buffer = zelf.buffer.write(); - *buffer = BufferedIO::new(Cursor::new(raw_bytes)); - let byte_pos = StringIO::char_offset_to_byte(buffer.cursor.get_ref(), pos as usize); - buffer - .seek(SeekFrom::Start(byte_pos as u64)) - .map_err(|err| os_err(vm, err))?; + *buffer = StringIOBuffer::new(raw_bytes); + buffer.seek(pos as usize); drop(buffer); zelf.newline.store(newline); let mut seennl = SeenNewline::empty(); diff --git a/extra_tests/snippets/stdlib_io.py b/extra_tests/snippets/stdlib_io.py index b0f497359..4aed3ebbf 100644 --- a/extra_tests/snippets/stdlib_io.py +++ b/extra_tests/snippets/stdlib_io.py @@ -3,6 +3,39 @@ from io import BufferedReader, BytesIO, FileIO, RawIOBase, StringIO, TextIOWrapp from testutils import assert_raises +# Overwriting replaces characters, even when their encoded widths differ. +for initial, replacement in ( + ("abc", "é"), + ("ébc", "a"), + ("а😀в", "XY"), + ("a\ud800c", "\udfffé"), +): + stream = StringIO("prefix:" + initial) + stream.seek(7) + assert stream.write(replacement) == len(replacement) + assert stream.getvalue() == "prefix:" + replacement + initial[len(replacement) :] + assert stream.tell() == 7 + len(replacement) + assert stream.read() == initial[len(replacement) :] + assert stream.tell() == 7 + len(initial) + +# Truncation keeps the character position, then pads in characters on writing. +stream = StringIO("é😀abc") +stream.seek(4) +stream.truncate(1) +assert stream.tell() == 4 +assert stream.read() == "" +assert stream.tell() == 4 +assert stream.write("") == 0 +assert stream.getvalue() == "é" +state = stream.__getstate__() +stream = StringIO() +stream.__setstate__(state) +assert stream.tell() == 4 +stream.write("Ж") +assert stream.getvalue() == "é\0\0\0Ж" +assert stream.tell() == 5 + + fi = FileIO("README.md") assert fi.seekable() bb = BufferedReader(fi)