diff --git a/newsfragments/5870.added.md b/newsfragments/5870.added.md new file mode 100644 index 00000000000..42907aaa792 --- /dev/null +++ b/newsfragments/5870.added.md @@ -0,0 +1 @@ +Add `obj()` getter to `PyUntypedBuffer` to retrieve the Python object owning the buffer. diff --git a/src/buffer.rs b/src/buffer.rs index 51f6563b690..4a46e003c30 100644 --- a/src/buffer.rs +++ b/src/buffer.rs @@ -570,6 +570,24 @@ impl PyUntypedBuffer { self.raw().buf } + ///Returns the Python object that owns the buffer data. + /// + ///This is the object that was passed to [`PyBuffer::get()`] when the buffer was created. + /// Returns the Python object that owns the buffer data. + /// + /// This is the object passed to [`PyUntypedBuffer::get()`]. + /// Calling this before [`release()`][Self::release] allows you to clone an owned reference + /// and keeps the object alive after the buffer is released. + pub fn obj<'py>(&self, py: Python<'py>) -> Option> { + let ptr = self.raw().obj; + // SAFETY: Py_buffer.obj is a borrowed reference to a Python object, so it is always valid to create a Bound from it + if ptr.is_null() { + None + } else { + Some(unsafe { Bound::from_borrowed_ptr(py, ptr) }) + } + } + /// Gets a pointer to the specified item. /// /// If `indices.len() < self.dimensions()`, returns the start address of the sub-array at the specified dimension. @@ -1044,4 +1062,27 @@ mod tests { assert_eq!(typed.shape(), [5]); }); } + + #[test] + fn test_obj_getter() { + Python::attach(|py| { + let bytes = PyBytes::new(py, b"hello"); + let buf = PyUntypedBuffer::get(bytes.as_any()).unwrap(); + + // obj() returns the same object that owns the buffer + let owner = buf.obj(py).unwrap(); + assert!(owner.is_instance_of::()); + assert!(owner.is(&bytes)); + + // can keep the owner alive after releasing the buffer + let owner_ref: crate::Py = owner.unbind(); + buf.release(py); + drop(bytes); + // owner_ref still valid after buffer and original are dropped + Python::attach(|py| { + let rebound = owner_ref.bind(py); + assert!(rebound.is_instance_of::()); + }); + }); + } }