Skip to content

Commit a63e3ac

Browse files
youknowoneyouknowone
committed
Rework the buffer protocol around managed exports and view offsets
Give `PyBuffer` the `_PyManagedBufferObject` shape: one `bf_getbuffer` acquisition is shared by every handle taken from it, cloning takes another share instead of re-acquiring, and the exporter's release runs once when the last share goes away. Remove `retain`, the unsafe `drop_without_release`, the three `impl Drop`s and the `ManuallyDrop` that stood in for this. Add `abort_acquisition` so a failed request does not run `bf_releasebuffer`. Move the view start into `BufferDescriptor::offset`, the `Py_buffer.buf` analogue, and drop the separate `start` fields on `PyMemoryView` and `PyBufferWrapper`. Slicing goes through `SaturatedSlice::adjust_indices_start`, which reproduces `PySlice_AdjustIndices` and keeps the adjusted start. Fix `zip_eq` to take its contiguous fast path only when both last dimensions are contiguous, and make `for_each_segment` and `zip_eq` handle zero-length and zero-dimensional views. Add `BufferDescriptor::projected` so a request without `PyBUF_ND`, `PyBUF_STRIDES` or `PyBUF_FORMAT` receives a correspondingly reduced descriptor, and reject a request without `PyBUF_INDIRECT` against an exporter that has suboffsets. Copy the source first in `memoryview` slice assignment when both sides reach the same root exporter. Hold the export across the resize in `bytearray.extend`, take `y*` in `marshal.loads`, stop probing the buffer protocol in `FsPath`, rewrite `ord` over the concrete string types, fold `array`'s buffer slot into one `slot_as_buffer`, take `w*`/`y*` in `_overlapped`, and thread the new `offset` field through the `_ctypes` descriptors. Assisted-by: Claude
1 parent a985ae4 commit a63e3ac

20 files changed

Lines changed: 832 additions & 347 deletions

File tree

‎crates/stdlib/src/array.rs‎

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ pub mod array {
2727
ArgBytesLike, ArgIntoFloat, ArgIterable, KwArgs, OptionalArg, PyComparisonValue,
2828
},
2929
protocol::{
30-
BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
31-
PyMappingMethods, PySequenceMethods,
30+
BufferDescriptor, BufferFlags, BufferMethods, BufferResizeGuard, PyBuffer,
31+
PyIterReturn, PyMappingMethods, PySequenceMethods,
3232
},
3333
sequence::{OptionalRangeArgs, SequenceExt, SequenceMutExt},
3434
sliceable::{
@@ -1291,22 +1291,42 @@ pub mod array {
12911291
}
12921292
}
12931293

1294+
impl PyArray {
1295+
fn buffer_desc(&self) -> BufferDescriptor {
1296+
let array = self.read();
1297+
BufferDescriptor::format(
1298+
array.len() * array.itemsize(),
1299+
false,
1300+
array.itemsize(),
1301+
array.typecode_str().into(),
1302+
)
1303+
}
1304+
}
1305+
12941306
impl AsBuffer for PyArray {
12951307
const RELEASE_BUFFER: bool = true;
12961308

1309+
// array_buffer_getbuf, which reports the type code only when the request
1310+
// asked for a format.
1311+
fn slot_as_buffer(
1312+
zelf: &PyObject,
1313+
flags: BufferFlags,
1314+
vm: &VirtualMachine,
1315+
) -> PyResult<PyBuffer> {
1316+
let zelf = zelf
1317+
.downcast_ref::<Self>()
1318+
.ok_or_else(|| vm.new_type_error("unexpected payload for as_buffer"))?;
1319+
let desc = zelf.buffer_desc().projected(flags);
1320+
flags.check_writable(desc.readonly, "Object is not writable.", vm)?;
1321+
Ok(PyBuffer::new(zelf.to_owned().into(), desc, &BUFFER_METHODS))
1322+
}
1323+
12971324
fn as_buffer(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
1298-
let array = zelf.read();
1299-
let buf = PyBuffer::new(
1325+
Ok(PyBuffer::new(
13001326
zelf.to_owned().into(),
1301-
BufferDescriptor::format(
1302-
array.len() * array.itemsize(),
1303-
false,
1304-
array.itemsize(),
1305-
array.typecode_str().into(),
1306-
),
1327+
zelf.buffer_desc(),
13071328
&BUFFER_METHODS,
1308-
);
1309-
Ok(buf)
1329+
))
13101330
}
13111331
}
13121332

‎crates/stdlib/src/overlapped.rs‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ mod _overlapped {
1212
builtins::{PyBaseExceptionRef, PyBytesRef, PyModule, PyStrRef, PyTupleRef, PyType},
1313
common::lock::PyMutex,
1414
convert::{ToPyException, ToPyObject},
15-
function::OptionalArg,
15+
function::{ArgBytesLike, ArgMemoryBuffer, OptionalArg},
1616
object::{Traverse, TraverseFn},
1717
protocol::PyBuffer,
1818
types::{Constructor, Destructor},
@@ -428,12 +428,14 @@ mod _overlapped {
428428
fn ReadFileInto(
429429
zelf: &Py<Self>,
430430
handle: isize,
431-
buf: PyBuffer,
431+
// w*, as _overlapped.Overlapped.ReadFileInto takes
432+
buf: ArgMemoryBuffer,
432433
vm: &VirtualMachine,
433434
) -> PyResult {
434435
use host_winapi::{
435436
ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS,
436437
};
438+
let buf: PyBuffer = buf.into();
437439

438440
let mut inner = zelf.inner.lock();
439441
if !matches!(inner.data, OverlappedData::None) {
@@ -530,13 +532,15 @@ mod _overlapped {
530532
fn WSARecvInto(
531533
zelf: &Py<Self>,
532534
handle: isize,
533-
buf: PyBuffer,
535+
// w*, as _overlapped.Overlapped.WSARecvInto takes
536+
buf: ArgMemoryBuffer,
534537
flags: u32,
535538
vm: &VirtualMachine,
536539
) -> PyResult {
537540
use host_winapi::{
538541
ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS,
539542
};
543+
let buf: PyBuffer = buf.into();
540544

541545
let mut inner = zelf.inner.lock();
542546
if !matches!(inner.data, OverlappedData::None) {
@@ -583,10 +587,12 @@ mod _overlapped {
583587
fn WriteFile(
584588
zelf: &Py<Self>,
585589
handle: isize,
586-
buf: PyBuffer,
590+
// y*, as _overlapped.Overlapped.WriteFile takes
591+
buf: ArgBytesLike,
587592
vm: &VirtualMachine,
588593
) -> PyResult {
589594
use host_winapi::{ERROR_IO_PENDING, ERROR_SUCCESS};
595+
let buf: PyBuffer = buf.into();
590596

591597
let mut inner = zelf.inner.lock();
592598
if !matches!(inner.data, OverlappedData::None) {
@@ -629,11 +635,13 @@ mod _overlapped {
629635
fn WSASend(
630636
zelf: &Py<Self>,
631637
handle: isize,
632-
buf: PyBuffer,
638+
// y*, as _overlapped.Overlapped.WSASend takes
639+
buf: ArgBytesLike,
633640
flags: u32,
634641
vm: &VirtualMachine,
635642
) -> PyResult {
636643
use host_winapi::{ERROR_IO_PENDING, ERROR_SUCCESS};
644+
let buf: PyBuffer = buf.into();
637645

638646
let mut inner = zelf.inner.lock();
639647
if !matches!(inner.data, OverlappedData::None) {
@@ -870,12 +878,14 @@ mod _overlapped {
870878
fn WSASendTo(
871879
zelf: &Py<Self>,
872880
handle: isize,
873-
buf: PyBuffer,
881+
// y*, as _overlapped.Overlapped.WSASendTo takes
882+
buf: ArgBytesLike,
874883
flags: u32,
875884
address: PyTupleRef,
876885
vm: &VirtualMachine,
877886
) -> PyResult {
878887
use host_winapi::{ERROR_IO_PENDING, ERROR_SUCCESS};
888+
let buf: PyBuffer = buf.into();
879889

880890
let mut inner = zelf.inner.lock();
881891
if !matches!(inner.data, OverlappedData::None) {
@@ -1001,14 +1011,16 @@ mod _overlapped {
10011011
fn WSARecvFromInto(
10021012
zelf: &Py<Self>,
10031013
handle: isize,
1004-
buf: PyBuffer,
1014+
// w*, as _overlapped.Overlapped.WSARecvFromInto takes
1015+
buf: ArgMemoryBuffer,
10051016
size: u32,
10061017
flags: OptionalArg<u32>,
10071018
vm: &VirtualMachine,
10081019
) -> PyResult {
10091020
use host_winapi::{
10101021
ERROR_BROKEN_PIPE, ERROR_IO_PENDING, ERROR_MORE_DATA, ERROR_SUCCESS,
10111022
};
1023+
let buf: PyBuffer = buf.into();
10121024

10131025
let mut inner = zelf.inner.lock();
10141026
if !matches!(inner.data, OverlappedData::None) {

‎crates/vm/src/anystr.rs‎

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use num_traits::{cast::ToPrimitive, sign::Signed};
44
use rustpython_unicode::case;
55

66
use crate::{
7-
Py, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
7+
AsObject, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
88
builtins::{PyIntRef, PyTuple},
99
convert::TryFromBorrowedObject,
1010
function::OptionalOption,
@@ -481,19 +481,25 @@ where
481481
F: Fn(T) -> PyResult<bool>,
482482
M: Fn(&PyObject) -> String,
483483
{
484-
if let Ok(single) = obj.try_to_value::<T>(vm) {
485-
(predicate)(single)
486-
} else {
487-
let tuple: &Py<PyTuple> = obj
488-
.try_to_value(vm)
489-
.map_err(|_| vm.new_type_error((message)(obj)))?;
490-
491-
for obj in tuple {
492-
if single_or_tuple_any(obj, predicate, message, vm)? {
484+
// _Py_bytes_tailmatch: a tuple is taken apart before anything is converted, and
485+
// each item is converted on its own terms, so a tuple of tuples is not an affix.
486+
if let Some(tuple) = obj.downcast_ref::<PyTuple>() {
487+
for item in tuple {
488+
if (predicate)(item.try_to_value::<T>(vm)?)? {
493489
return Ok(true);
494490
}
495491
}
496-
497-
Ok(false)
492+
return Ok(false);
498493
}
494+
495+
// Only the argument simply being the wrong kind of object is reported as such;
496+
// whatever the conversion itself raised belongs to the caller.
497+
let single = obj.try_to_value::<T>(vm).map_err(|exc| {
498+
if exc.fast_isinstance(vm.ctx.exceptions.type_error) {
499+
vm.new_type_error((message)(obj))
500+
} else {
501+
exc
502+
}
503+
})?;
504+
(predicate)(single)
499505
}

‎crates/vm/src/builtins/bytearray.rs‎

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use crate::{
88
VirtualMachine,
99
anystr::{self, AnyStr},
1010
atomic_func,
11-
byte::{bytes_from_object, bytes_from_setslice_value, value_from_object},
11+
byte::{bytes_from_object, value_from_object},
1212
bytes_inner::{
1313
ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
1414
ByteInnerSplitOptions, ByteInnerSub, ByteInnerTranslateOptions, DecodeArgs, PyBytesInner,
@@ -611,12 +611,34 @@ impl Py<PyByteArray> {
611611
#[pymethod]
612612
fn extend(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
613613
if self.is(&object) {
614-
PyByteArray::irepeat(self, 2, vm)
615-
} else {
616-
let items = bytes_from_setslice_value(vm, &object)?;
617-
self.try_resizable(vm)?.elements.extend(items);
618-
Ok(())
614+
return PyByteArray::irepeat(self, 2, vm);
619615
}
616+
// bytearray_setslice keeps the export alive across the resize, so a value
617+
// looking at this bytearray is what stops it from growing.
618+
let buffer = object
619+
.check_buffer()
620+
.then(|| {
621+
PyBuffer::from_object(vm, &object, BufferFlags::SIMPLE).map_err(|_| {
622+
// What an exporter refuses to hand out leaves the value simply
623+
// not usable here, whatever the exporter's own complaint was.
624+
vm.new_type_error(format!(
625+
"can't set bytearray slice from {}",
626+
object.class().name()
627+
))
628+
})
629+
})
630+
.transpose()?;
631+
let items = match &buffer {
632+
Some(buffer) => buffer
633+
.as_contiguous()
634+
.ok_or_else(|| {
635+
vm.new_buffer_error("non-contiguous buffer is not a bytes-like object")
636+
})?
637+
.to_vec(),
638+
None => bytes_from_object(vm, &object)?,
639+
};
640+
self.try_resizable(vm)?.elements.extend(items);
641+
Ok(())
620642
}
621643

622644
#[pymethod]

0 commit comments

Comments
 (0)