// TODO https://github.com/PyO3/pyo3/issues/5487 #![allow(clippy::undocumented_unsafe_blocks)] #![cfg(feature = "macros")] #![cfg(any(not(Py_LIMITED_API), Py_3_11))] use pyo3::buffer::PyBuffer; use pyo3::exceptions::PyBufferError; use pyo3::ffi; use pyo3::prelude::*; use pyo3::types::IntoPyDict; use std::ffi::CString; use std::ffi::{c_int, c_void}; use std::ptr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; mod test_utils; #[pyclass] struct TestBufferClass { vec: Vec, drop_called: Arc, } #[pymethods] impl TestBufferClass { unsafe fn __getbuffer__( slf: Bound<'_, Self>, view: *mut ffi::Py_buffer, flags: c_int, ) -> PyResult<()> { unsafe { fill_view_from_readonly_data(view, flags, &slf.borrow().vec, slf.into_any()) } } unsafe fn __releasebuffer__(&self, view: *mut ffi::Py_buffer) { // Release memory held by the format string drop(unsafe { CString::from_raw((*view).format) }); } } impl Drop for TestBufferClass { fn drop(&mut self) { print!("dropped"); self.drop_called.store(true, Ordering::Relaxed); } } #[test] fn test_buffer() { let drop_called = Arc::new(AtomicBool::new(false)); Python::attach(|py| { let instance = Py::new( py, TestBufferClass { vec: vec![b' ', b'2', b'3'], drop_called: drop_called.clone(), }, ) .unwrap(); let env = [("ob", instance)].into_py_dict(py).unwrap(); py_assert!(py, *env, "bytes(ob) == b' 23'"); }); assert!(drop_called.load(Ordering::Relaxed)); } #[test] fn test_buffer_referenced() { let drop_called = Arc::new(AtomicBool::new(false)); let buf = { let input = vec![b' ', b'2', b'3']; Python::attach(|py| { let instance = TestBufferClass { vec: input.clone(), drop_called: drop_called.clone(), } .into_pyobject(py) .unwrap(); let buf = PyBuffer::::get(&instance).unwrap(); assert_eq!(buf.to_vec(py).unwrap(), input); drop(instance); buf }) }; assert!(!drop_called.load(Ordering::Relaxed)); Python::attach(|_| { drop(buf); }); assert!(drop_called.load(Ordering::Relaxed)); } #[test] fn test_releasebuffer_unraisable_error() { use pyo3::exceptions::PyValueError; use test_utils::UnraisableCapture; #[pyclass] struct ReleaseBufferError {} #[pymethods] impl ReleaseBufferError { unsafe fn __getbuffer__( slf: Bound<'_, Self>, view: *mut ffi::Py_buffer, flags: c_int, ) -> PyResult<()> { static BUF_BYTES: &[u8] = b"hello world"; unsafe { fill_view_from_readonly_data(view, flags, BUF_BYTES, slf.into_any()) } } unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) -> PyResult<()> { Err(PyValueError::new_err("oh dear")) } } Python::attach(|py| { let instance = Py::new(py, ReleaseBufferError {}).unwrap(); let (err, object) = UnraisableCapture::enter(py, |capture| { let env = [("ob", instance.clone_ref(py))].into_py_dict(py).unwrap(); assert!(capture.take_capture().is_none()); py_assert!(py, *env, "bytes(ob) == b'hello world'"); capture.take_capture().unwrap() }); assert_eq!(err.to_string(), "ValueError: oh dear"); assert!(object.is(&instance)); }); } /// # Safety /// /// `view` must be a valid pointer to ffi::Py_buffer, or null /// `data` must outlive the Python lifetime of `owner` (i.e. data must be owned by owner, or data /// must be static data) unsafe fn fill_view_from_readonly_data( view: *mut ffi::Py_buffer, flags: c_int, data: &[u8], owner: Bound<'_, PyAny>, ) -> PyResult<()> { if view.is_null() { return Err(PyBufferError::new_err("View is null")); } if (flags & ffi::PyBUF_WRITABLE) == ffi::PyBUF_WRITABLE { return Err(PyBufferError::new_err("Object is not writable")); } unsafe { (*view).obj = owner.into_ptr(); (*view).buf = data.as_ptr() as *mut c_void; (*view).len = data.len() as isize; (*view).readonly = 1; (*view).itemsize = 1; (*view).format = if (flags & ffi::PyBUF_FORMAT) == ffi::PyBUF_FORMAT { let msg = CString::new("B").unwrap(); msg.into_raw() } else { ptr::null_mut() }; (*view).ndim = 1; (*view).shape = if (flags & ffi::PyBUF_ND) == ffi::PyBUF_ND { &mut (*view).len } else { ptr::null_mut() }; (*view).strides = if (flags & ffi::PyBUF_STRIDES) == ffi::PyBUF_STRIDES { &mut (*view).itemsize } else { ptr::null_mut() }; (*view).suboffsets = ptr::null_mut(); (*view).internal = ptr::null_mut(); } Ok(()) }