Skip to content

Commit f26e909

Browse files
committed
Release BytesIO and StringIO buffers on close
Assisted-by: Codex:GPT-6
1 parent eb2ffc6 commit f26e909

2 files changed

Lines changed: 98 additions & 8 deletions

File tree

‎crates/vm/src/stdlib/_io.rs‎

Lines changed: 64 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4895,21 +4895,25 @@ mod _io {
48954895
let raw_bytes = initial_value.as_ref().map_or_else(Vec::new, |v| {
48964896
Self::translate_newlines(v.as_wtf8(), newline).into_bytes()
48974897
});
4898-
*zelf.buffer.write() = StringIOBuffer::new(raw_bytes);
4898+
let mut buffer = zelf.buffer.write();
4899+
*buffer = StringIOBuffer::new(raw_bytes);
48994900
zelf.newline.store(newline);
49004901
zelf.seennl.store(SeenNewline::empty());
4901-
if let Some(initial_value) = initial_value {
4902+
if let Some(initial_value) = initial_value.as_ref() {
49024903
zelf.observe_newlines(initial_value.as_wtf8(), newline);
49034904
}
4905+
zelf.closed.store(false);
49044906
Ok(())
49054907
}
49064908
}
49074909

49084910
impl StringIO {
49094911
fn buffer(&self, vm: &VirtualMachine) -> PyResult<PyRwLockWriteGuard<'_, StringIOBuffer>> {
4912+
let buffer = self.buffer.write();
49104913
if !self.closed.load() {
4911-
Ok(self.buffer.write())
4914+
Ok(buffer)
49124915
} else {
4916+
drop(buffer);
49134917
Err(io_closed_error(vm))
49144918
}
49154919
}
@@ -4983,7 +4987,9 @@ mod _io {
49834987

49844988
#[pymethod]
49854989
fn close(&self) {
4990+
let mut buffer = self.buffer.write();
49864991
self.closed.store(true);
4992+
*buffer = StringIOBuffer::new(Vec::new());
49874993
}
49884994

49894995
// Replace characters in the underlying vector, regardless of their byte widths.
@@ -5141,7 +5147,7 @@ mod _io {
51415147

51425148
// Set content and position
51435149
let raw_bytes = content.as_bytes().to_vec();
5144-
let mut buffer = zelf.buffer.write();
5150+
let mut buffer = zelf.buffer(vm)?;
51455151
*buffer = StringIOBuffer::new(raw_bytes);
51465152
buffer.seek(pos as usize);
51475153
drop(buffer);
@@ -5209,17 +5215,22 @@ mod _io {
52095215

52105216
let raw_bytes = args
52115217
.initial_bytes
5218+
.as_ref()
52125219
.map_or_else(Vec::new, |input| input.borrow_buf().to_vec());
5213-
*zelf.buffer.write() = BufferedIO::new(Cursor::new(raw_bytes));
5220+
let mut buffer = zelf.try_resizable(vm)?;
5221+
*buffer = BufferedIO::new(Cursor::new(raw_bytes));
5222+
zelf.closed.store(false);
52145223
Ok(())
52155224
}
52165225
}
52175226

52185227
impl BytesIO {
52195228
fn buffer(&self, vm: &VirtualMachine) -> PyResult<PyRwLockWriteGuard<'_, BufferedIO>> {
5229+
let buffer = self.buffer.write();
52205230
if !self.closed.load() {
5221-
Ok(self.buffer.write())
5231+
Ok(buffer)
52225232
} else {
5233+
drop(buffer);
52235234
Err(io_closed_error(vm))
52245235
}
52255236
}
@@ -5340,6 +5351,10 @@ mod _io {
53405351
return Err(io_closed_error(vm));
53415352
}
53425353
let mut buffer = self.try_resizable(vm)?;
5354+
if self.closed.load() {
5355+
drop(buffer);
5356+
return Err(io_closed_error(vm));
5357+
}
53435358
let pos = pos.try_usize(vm)?;
53445359
Ok(buffer.truncate(pos))
53455360
}
@@ -5356,7 +5371,15 @@ mod _io {
53565371
vm.new_buffer_error("Existing exports of data: object cannot be closed")
53575372
);
53585373
}
5374+
let mut buffer = self.buffer.write();
5375+
if self.exports.load() > 0 {
5376+
drop(buffer);
5377+
return Err(
5378+
vm.new_buffer_error("Existing exports of data: object cannot be closed")
5379+
);
5380+
}
53595381
self.closed.store(true);
5382+
*buffer = BufferedIO::new(Cursor::new(Vec::new()));
53605383
Ok(())
53615384
}
53625385

@@ -5403,6 +5426,10 @@ mod _io {
54035426

54045427
// Check exports and set content (like CHECK_EXPORTS)
54055428
let mut buffer = zelf.try_resizable(vm)?;
5429+
if zelf.closed.load() {
5430+
drop(buffer);
5431+
return Err(vm.new_value_error("__setstate__ on closed file"));
5432+
}
54065433
*buffer = BufferedIO::new(Cursor::new(content.as_bytes().to_vec()));
54075434
buffer
54085435
.seek(SeekFrom::Start(pos))
@@ -5437,15 +5464,18 @@ mod _io {
54375464
impl PyRef<BytesIO> {
54385465
#[pymethod]
54395466
fn getbuffer(self, vm: &VirtualMachine) -> PyResult<PyMemoryView> {
5467+
let storage = self.buffer.read();
54405468
if self.closed.load() {
5469+
drop(storage);
54415470
return Err(vm.new_value_error("I/O operation on closed file."));
54425471
}
5443-
let len = self.buffer.read().cursor.get_ref().len();
5472+
let len = storage.cursor.get_ref().len();
54445473
let buffer = PyBuffer::new(
5445-
self.into(),
5474+
self.clone().into(),
54465475
BufferDescriptor::simple(len, false),
54475476
&BYTES_IO_BUFFER_METHODS,
54485477
);
5478+
drop(storage);
54495479
let view = PyMemoryView::from_buffer(buffer, vm)?;
54505480
Ok(view)
54515481
}
@@ -5919,6 +5949,32 @@ mod _io {
59195949
mod tests {
59205950
use super::*;
59215951

5952+
#[test]
5953+
fn closed_memory_streams_release_storage() {
5954+
crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
5955+
let text = StringIO {
5956+
_base: Default::default(),
5957+
buffer: PyRwLock::new(StringIOBuffer::new(vec![b'x'; 1024])),
5958+
newline: AtomicCell::new(Newlines::Lf),
5959+
seennl: AtomicCell::new(SeenNewline::empty()),
5960+
closed: AtomicCell::new(false),
5961+
};
5962+
let text = text.into_ref(&vm.ctx);
5963+
text.close();
5964+
assert_eq!(text.buffer.read().bytes.cursor.get_ref().capacity(), 0);
5965+
5966+
let bytes = BytesIO {
5967+
_base: Default::default(),
5968+
buffer: PyRwLock::new(BufferedIO::new(Cursor::new(vec![0; 1024]))),
5969+
closed: AtomicCell::new(false),
5970+
exports: AtomicCell::new(0),
5971+
};
5972+
let bytes = bytes.into_ref(&vm.ctx);
5973+
bytes.close(vm).unwrap();
5974+
assert_eq!(bytes.buffer.read().cursor.get_ref().capacity(), 0);
5975+
});
5976+
}
5977+
59225978
#[test]
59235979
fn buffered_read() {
59245980
let data = vec![1, 2, 3, 4];

‎extra_tests/snippets/stdlib_io.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,3 +288,37 @@ def _text_cookie(
288288
else:
289289
assert _textio.read(50) is not None
290290
_textio.tell()
291+
292+
293+
# Reinitialization releases arguments after unlocking the stream.
294+
for _type, _content in ((BytesIO, b"hello"), (StringIO, "hello")):
295+
_stream = _type(_content)
296+
_stream.close()
297+
_stream.__init__(_content)
298+
assert not _stream.closed
299+
assert _stream.getvalue() == _content
300+
301+
class ClosingValue(type(_content)):
302+
def __del__(self):
303+
_stream.close()
304+
305+
_stream.__init__(ClosingValue(_content))
306+
assert _stream.closed
307+
308+
309+
_stream = BytesIO(b"old")
310+
_released_values = []
311+
312+
313+
class ClosingBuffer:
314+
def __buffer__(self, flags):
315+
return memoryview(b"new")
316+
317+
def __release_buffer__(self, view):
318+
_released_values.append(_stream.getvalue())
319+
_stream.close()
320+
321+
322+
_stream.__init__(ClosingBuffer())
323+
assert _released_values == [b"new"]
324+
assert _stream.closed

0 commit comments

Comments
 (0)