#![cfg(any(not(Py_LIMITED_API), Py_3_11))]
use crate::{
err, exceptions::PyBufferError, ffi, AsPyPointer, FromPyObject, PyAny, PyResult, Python,
};
use std::marker::PhantomData;
use std::os::raw;
use std::pin::Pin;
use std::{cell, mem, ptr, slice};
use std::{ffi::CStr, fmt::Debug};
#[repr(transparent)]
pub struct PyBuffer<T>(Pin<Box<ffi::Py_buffer>>, PhantomData<T>);
unsafe impl<T> Send for PyBuffer<T> {}
unsafe impl<T> Sync for PyBuffer<T> {}
impl<T> Debug for PyBuffer<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PyBuffer")
.field("buf", &self.0.buf)
.field("obj", &self.0.obj)
.field("len", &self.0.len)
.field("itemsize", &self.0.itemsize)
.field("readonly", &self.0.readonly)
.field("ndim", &self.0.ndim)
.field("format", &self.0.format)
.field("shape", &self.0.shape)
.field("strides", &self.0.strides)
.field("suboffsets", &self.0.suboffsets)
.field("internal", &self.0.internal)
.finish()
}
}
#[derive(Copy, Clone, Eq, PartialEq)]
pub enum ElementType {
SignedInteger { bytes: usize },
UnsignedInteger { bytes: usize },
Bool,
Float { bytes: usize },
Unknown,
}
impl ElementType {
pub fn from_format(format: &CStr) -> ElementType {
match format.to_bytes() {
[char] | [b'@', char] => native_element_type_from_type_char(*char),
[modifier, char] if matches!(modifier, b'=' | b'<' | b'>' | b'!') => {
standard_element_type_from_type_char(*char)
}
_ => ElementType::Unknown,
}
}
}
fn native_element_type_from_type_char(type_char: u8) -> ElementType {
use self::ElementType::*;
match type_char {
b'c' => UnsignedInteger {
bytes: mem::size_of::<raw::c_char>(),
},
b'b' => SignedInteger {
bytes: mem::size_of::<raw::c_schar>(),
},
b'B' => UnsignedInteger {
bytes: mem::size_of::<raw::c_uchar>(),
},
b'?' => Bool,
b'h' => SignedInteger {
bytes: mem::size_of::<raw::c_short>(),
},
b'H' => UnsignedInteger {
bytes: mem::size_of::<raw::c_ushort>(),
},
b'i' => SignedInteger {
bytes: mem::size_of::<raw::c_int>(),
},
b'I' => UnsignedInteger {
bytes: mem::size_of::<raw::c_uint>(),
},
b'l' => SignedInteger {
bytes: mem::size_of::<raw::c_long>(),
},
b'L' => UnsignedInteger {
bytes: mem::size_of::<raw::c_ulong>(),
},
b'q' => SignedInteger {
bytes: mem::size_of::<raw::c_longlong>(),
},
b'Q' => UnsignedInteger {
bytes: mem::size_of::<raw::c_ulonglong>(),
},
b'n' => SignedInteger {
bytes: mem::size_of::<libc::ssize_t>(),
},
b'N' => UnsignedInteger {
bytes: mem::size_of::<libc::size_t>(),
},
b'e' => Float { bytes: 2 },
b'f' => Float { bytes: 4 },
b'd' => Float { bytes: 8 },
_ => Unknown,
}
}
fn standard_element_type_from_type_char(type_char: u8) -> ElementType {
use self::ElementType::*;
match type_char {
b'c' | b'B' => UnsignedInteger { bytes: 1 },
b'b' => SignedInteger { bytes: 1 },
b'?' => Bool,
b'h' => SignedInteger { bytes: 2 },
b'H' => UnsignedInteger { bytes: 2 },
b'i' | b'l' => SignedInteger { bytes: 4 },
b'I' | b'L' => UnsignedInteger { bytes: 4 },
b'q' => SignedInteger { bytes: 8 },
b'Q' => UnsignedInteger { bytes: 8 },
b'e' => Float { bytes: 2 },
b'f' => Float { bytes: 4 },
b'd' => Float { bytes: 8 },
_ => Unknown,
}
}
#[cfg(target_endian = "little")]
fn is_matching_endian(c: u8) -> bool {
c == b'@' || c == b'=' || c == b'>'
}
#[cfg(target_endian = "big")]
fn is_matching_endian(c: u8) -> bool {
c == b'@' || c == b'=' || c == b'>' || c == b'!'
}
pub unsafe trait Element: Copy {
fn is_compatible_format(format: &CStr) -> bool;
}
impl<'source, T: Element> FromPyObject<'source> for PyBuffer<T> {
fn extract(obj: &PyAny) -> PyResult<PyBuffer<T>> {
Self::get(obj)
}
}
impl<T: Element> PyBuffer<T> {
pub fn get(obj: &PyAny) -> PyResult<PyBuffer<T>> {
let mut buf = Box::new(mem::MaybeUninit::uninit());
let buf: Box<ffi::Py_buffer> = unsafe {
err::error_on_minusone(
obj.py(),
ffi::PyObject_GetBuffer(obj.as_ptr(), buf.as_mut_ptr(), ffi::PyBUF_FULL_RO),
)?;
mem::transmute(buf)
};
let buf = PyBuffer(Pin::from(buf), PhantomData);
if buf.0.shape.is_null() {
Err(PyBufferError::new_err("shape is null"))
} else if buf.0.strides.is_null() {
Err(PyBufferError::new_err("strides is null"))
} else if mem::size_of::<T>() != buf.item_size() || !T::is_compatible_format(buf.format()) {
Err(PyBufferError::new_err(format!(
"buffer contents are not compatible with {}",
std::any::type_name::<T>()
)))
} else if buf.0.buf.align_offset(mem::align_of::<T>()) != 0 {
Err(PyBufferError::new_err(format!(
"buffer contents are insufficiently aligned for {}",
std::any::type_name::<T>()
)))
} else {
Ok(buf)
}
}
#[inline]
pub fn buf_ptr(&self) -> *mut raw::c_void {
self.0.buf
}
pub fn get_ptr(&self, indices: &[usize]) -> *mut raw::c_void {
let shape = &self.shape()[..indices.len()];
for i in 0..indices.len() {
assert!(indices[i] < shape[i]);
}
unsafe {
ffi::PyBuffer_GetPointer(
#[cfg(Py_3_11)]
&*self.0,
#[cfg(not(Py_3_11))]
{
&*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
},
#[cfg(Py_3_11)]
{
indices.as_ptr() as *const ffi::Py_ssize_t
},
#[cfg(not(Py_3_11))]
{
indices.as_ptr() as *mut ffi::Py_ssize_t
},
)
}
}
#[inline]
pub fn readonly(&self) -> bool {
self.0.readonly != 0
}
#[inline]
pub fn item_size(&self) -> usize {
self.0.itemsize as usize
}
#[inline]
pub fn item_count(&self) -> usize {
(self.0.len as usize) / (self.0.itemsize as usize)
}
#[inline]
pub fn len_bytes(&self) -> usize {
self.0.len as usize
}
#[inline]
pub fn dimensions(&self) -> usize {
self.0.ndim as usize
}
#[inline]
pub fn shape(&self) -> &[usize] {
unsafe { slice::from_raw_parts(self.0.shape as *const usize, self.0.ndim as usize) }
}
#[inline]
pub fn strides(&self) -> &[isize] {
unsafe { slice::from_raw_parts(self.0.strides, self.0.ndim as usize) }
}
#[inline]
pub fn suboffsets(&self) -> Option<&[isize]> {
unsafe {
if self.0.suboffsets.is_null() {
None
} else {
Some(slice::from_raw_parts(
self.0.suboffsets,
self.0.ndim as usize,
))
}
}
}
#[inline]
pub fn format(&self) -> &CStr {
if self.0.format.is_null() {
CStr::from_bytes_with_nul(b"B\0").unwrap()
} else {
unsafe { CStr::from_ptr(self.0.format) }
}
}
#[inline]
pub fn is_c_contiguous(&self) -> bool {
unsafe {
ffi::PyBuffer_IsContiguous(
&*self.0 as *const ffi::Py_buffer,
b'C' as std::os::raw::c_char,
) != 0
}
}
#[inline]
pub fn is_fortran_contiguous(&self) -> bool {
unsafe {
ffi::PyBuffer_IsContiguous(
&*self.0 as *const ffi::Py_buffer,
b'F' as std::os::raw::c_char,
) != 0
}
}
pub fn as_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [ReadOnlyCell<T>]> {
if self.is_c_contiguous() {
unsafe {
Some(slice::from_raw_parts(
self.0.buf as *mut ReadOnlyCell<T>,
self.item_count(),
))
}
} else {
None
}
}
pub fn as_mut_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [cell::Cell<T>]> {
if !self.readonly() && self.is_c_contiguous() {
unsafe {
Some(slice::from_raw_parts(
self.0.buf as *mut cell::Cell<T>,
self.item_count(),
))
}
} else {
None
}
}
pub fn as_fortran_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [ReadOnlyCell<T>]> {
if mem::size_of::<T>() == self.item_size() && self.is_fortran_contiguous() {
unsafe {
Some(slice::from_raw_parts(
self.0.buf as *mut ReadOnlyCell<T>,
self.item_count(),
))
}
} else {
None
}
}
pub fn as_fortran_mut_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [cell::Cell<T>]> {
if !self.readonly() && self.is_fortran_contiguous() {
unsafe {
Some(slice::from_raw_parts(
self.0.buf as *mut cell::Cell<T>,
self.item_count(),
))
}
} else {
None
}
}
pub fn copy_to_slice(&self, py: Python<'_>, target: &mut [T]) -> PyResult<()> {
self.copy_to_slice_impl(py, target, b'C')
}
pub fn copy_to_fortran_slice(&self, py: Python<'_>, target: &mut [T]) -> PyResult<()> {
self.copy_to_slice_impl(py, target, b'F')
}
fn copy_to_slice_impl(&self, py: Python<'_>, target: &mut [T], fort: u8) -> PyResult<()> {
if mem::size_of_val(target) != self.len_bytes() {
return Err(PyBufferError::new_err(format!(
"slice to copy to (of length {}) does not match buffer length of {}",
target.len(),
self.item_count()
)));
}
unsafe {
err::error_on_minusone(
py,
ffi::PyBuffer_ToContiguous(
target.as_ptr() as *mut raw::c_void,
#[cfg(Py_3_11)]
&*self.0,
#[cfg(not(Py_3_11))]
{
&*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
},
self.0.len,
fort as std::os::raw::c_char,
),
)
}
}
pub fn to_vec(&self, py: Python<'_>) -> PyResult<Vec<T>> {
self.to_vec_impl(py, b'C')
}
pub fn to_fortran_vec(&self, py: Python<'_>) -> PyResult<Vec<T>> {
self.to_vec_impl(py, b'F')
}
fn to_vec_impl(&self, py: Python<'_>, fort: u8) -> PyResult<Vec<T>> {
let item_count = self.item_count();
let mut vec: Vec<T> = Vec::with_capacity(item_count);
unsafe {
err::error_on_minusone(
py,
ffi::PyBuffer_ToContiguous(
vec.as_ptr() as *mut raw::c_void,
#[cfg(Py_3_11)]
&*self.0,
#[cfg(not(Py_3_11))]
{
&*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
},
self.0.len,
fort as std::os::raw::c_char,
),
)?;
vec.set_len(item_count);
}
Ok(vec)
}
pub fn copy_from_slice(&self, py: Python<'_>, source: &[T]) -> PyResult<()> {
self.copy_from_slice_impl(py, source, b'C')
}
pub fn copy_from_fortran_slice(&self, py: Python<'_>, source: &[T]) -> PyResult<()> {
self.copy_from_slice_impl(py, source, b'F')
}
fn copy_from_slice_impl(&self, py: Python<'_>, source: &[T], fort: u8) -> PyResult<()> {
if self.readonly() {
return Err(PyBufferError::new_err("cannot write to read-only buffer"));
} else if mem::size_of_val(source) != self.len_bytes() {
return Err(PyBufferError::new_err(format!(
"slice to copy from (of length {}) does not match buffer length of {}",
source.len(),
self.item_count()
)));
}
unsafe {
err::error_on_minusone(
py,
ffi::PyBuffer_FromContiguous(
#[cfg(Py_3_11)]
&*self.0,
#[cfg(not(Py_3_11))]
{
&*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
},
#[cfg(Py_3_11)]
{
source.as_ptr() as *const raw::c_void
},
#[cfg(not(Py_3_11))]
{
source.as_ptr() as *mut raw::c_void
},
self.0.len,
fort as std::os::raw::c_char,
),
)
}
}
pub fn release(self, _py: Python<'_>) {
let mut mdself = mem::ManuallyDrop::new(self);
unsafe {
ffi::PyBuffer_Release(&mut *mdself.0);
let inner: *mut Pin<Box<ffi::Py_buffer>> = &mut mdself.0;
ptr::drop_in_place(inner);
}
}
}
impl<T> Drop for PyBuffer<T> {
fn drop(&mut self) {
Python::with_gil(|_| unsafe { ffi::PyBuffer_Release(&mut *self.0) });
}
}
#[repr(transparent)]
pub struct ReadOnlyCell<T: Element>(cell::UnsafeCell<T>);
impl<T: Element> ReadOnlyCell<T> {
#[inline]
pub fn get(&self) -> T {
unsafe { *self.0.get() }
}
#[inline]
pub fn as_ptr(&self) -> *const T {
self.0.get()
}
}
macro_rules! impl_element(
($t:ty, $f:ident) => {
unsafe impl Element for $t {
fn is_compatible_format(format: &CStr) -> bool {
let slice = format.to_bytes();
if slice.len() > 1 && !is_matching_endian(slice[0]) {
return false;
}
ElementType::from_format(format) == ElementType::$f { bytes: mem::size_of::<$t>() }
}
}
}
);
impl_element!(u8, UnsignedInteger);
impl_element!(u16, UnsignedInteger);
impl_element!(u32, UnsignedInteger);
impl_element!(u64, UnsignedInteger);
impl_element!(usize, UnsignedInteger);
impl_element!(i8, SignedInteger);
impl_element!(i16, SignedInteger);
impl_element!(i32, SignedInteger);
impl_element!(i64, SignedInteger);
impl_element!(isize, SignedInteger);
impl_element!(f32, Float);
impl_element!(f64, Float);
#[cfg(test)]
mod tests {
use super::PyBuffer;
use crate::ffi;
use crate::Python;
#[test]
fn test_compatible_size() {
assert_eq!(
std::mem::size_of::<ffi::Py_ssize_t>(),
std::mem::size_of::<usize>()
);
}
#[test]
fn test_bytes_buffer() {
Python::with_gil(|py| {
let bytes = py.eval("b'abcde'", None, None).unwrap();
let buffer = PyBuffer::get(bytes).unwrap();
assert_eq!(buffer.dimensions(), 1);
assert_eq!(buffer.item_count(), 5);
assert_eq!(buffer.format().to_str().unwrap(), "B");
assert_eq!(buffer.shape(), [5]);
assert!(buffer.is_c_contiguous());
assert!(buffer.is_fortran_contiguous());
let slice = buffer.as_slice(py).unwrap();
assert_eq!(slice.len(), 5);
assert_eq!(slice[0].get(), b'a');
assert_eq!(slice[2].get(), b'c');
assert_eq!(unsafe { *(buffer.get_ptr(&[1]) as *mut u8) }, b'b');
assert!(buffer.copy_to_slice(py, &mut [0u8]).is_err());
let mut arr = [0; 5];
buffer.copy_to_slice(py, &mut arr).unwrap();
assert_eq!(arr, b"abcde" as &[u8]);
assert!(buffer.copy_from_slice(py, &[0u8; 5]).is_err());
assert_eq!(buffer.to_vec(py).unwrap(), b"abcde");
});
}
#[test]
fn test_array_buffer() {
Python::with_gil(|py| {
let array = py
.import("array")
.unwrap()
.call_method("array", ("f", (1.0, 1.5, 2.0, 2.5)), None)
.unwrap();
let buffer = PyBuffer::get(array).unwrap();
assert_eq!(buffer.dimensions(), 1);
assert_eq!(buffer.item_count(), 4);
assert_eq!(buffer.format().to_str().unwrap(), "f");
assert_eq!(buffer.shape(), [4]);
let slice = buffer.as_slice(py).unwrap();
assert_eq!(slice.len(), 4);
assert_eq!(slice[0].get(), 1.0);
assert_eq!(slice[3].get(), 2.5);
let mut_slice = buffer.as_mut_slice(py).unwrap();
assert_eq!(mut_slice.len(), 4);
assert_eq!(mut_slice[0].get(), 1.0);
mut_slice[3].set(2.75);
assert_eq!(slice[3].get(), 2.75);
buffer
.copy_from_slice(py, &[10.0f32, 11.0, 12.0, 13.0])
.unwrap();
assert_eq!(slice[2].get(), 12.0);
assert_eq!(buffer.to_vec(py).unwrap(), [10.0, 11.0, 12.0, 13.0]);
});
}
}