pyo3/
buffer.rs

1#![cfg(any(not(Py_LIMITED_API), Py_3_11))]
2// Copyright (c) 2017 Daniel Grunwald
3//
4// Permission is hereby granted, free of charge, to any person obtaining a copy of this
5// software and associated documentation files (the "Software"), to deal in the Software
6// without restriction, including without limitation the rights to use, copy, modify, merge,
7// publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons
8// to whom the Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in all copies or
11// substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
14// INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
15// PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
16// FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
17// OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
18// DEALINGS IN THE SOFTWARE.
19
20//! `PyBuffer` implementation
21#[cfg(feature = "experimental-inspect")]
22use crate::inspect::{type_hint_identifier, PyStaticExpr};
23use crate::{err, exceptions::PyBufferError, ffi, FromPyObject, PyAny, PyResult, Python};
24use crate::{Borrowed, Bound, PyErr};
25use std::ffi::{
26    c_char, c_int, c_long, c_longlong, c_schar, c_short, c_uchar, c_uint, c_ulong, c_ulonglong,
27    c_ushort, c_void,
28};
29use std::marker::{PhantomData, PhantomPinned};
30use std::pin::Pin;
31use std::ptr::NonNull;
32use std::{cell, mem, ptr, slice};
33use std::{ffi::CStr, fmt::Debug};
34
35/// A typed form of [`PyUntypedBuffer`].
36#[repr(transparent)]
37pub struct PyBuffer<T>(PyUntypedBuffer, PhantomData<[T]>);
38
39/// Allows access to the underlying buffer used by a python object such as `bytes`, `bytearray` or `array.array`.
40#[repr(transparent)]
41pub struct PyUntypedBuffer(
42    // It is common for exporters filling `Py_buffer` struct to make it self-referential, e.g. see
43    // implementation of
44    // [`PyBuffer_FillInfo`](https://github.com/python/cpython/blob/2fd43a1ffe4ff1f6c46f6045bc327d6085c40fbf/Objects/abstract.c#L798-L802).
45    //
46    // Therefore we use `Pin<Box<...>>` to document for ourselves that the memory address of the `Py_buffer` is expected to be stable
47    Pin<Box<RawBuffer>>,
48);
49
50/// Wrapper around `ffi::Py_buffer` to be `!Unpin`.
51#[repr(transparent)]
52struct RawBuffer(ffi::Py_buffer, PhantomPinned);
53
54// PyBuffer send & sync guarantees are upheld by Python.
55unsafe impl Send for PyUntypedBuffer {}
56unsafe impl Sync for PyUntypedBuffer {}
57
58impl<T> Debug for PyBuffer<T> {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        debug_buffer("PyBuffer", &self.0, f)
61    }
62}
63
64impl Debug for PyUntypedBuffer {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        debug_buffer("PyUntypedBuffer", self, f)
67    }
68}
69
70fn debug_buffer(
71    name: &str,
72    b: &PyUntypedBuffer,
73    f: &mut std::fmt::Formatter<'_>,
74) -> std::fmt::Result {
75    let raw = b.raw();
76    f.debug_struct(name)
77        .field("buf", &raw.buf)
78        .field("obj", &raw.obj)
79        .field("len", &raw.len)
80        .field("itemsize", &raw.itemsize)
81        .field("readonly", &raw.readonly)
82        .field("ndim", &raw.ndim)
83        .field("format", &b.format())
84        .field("shape", &b.shape())
85        .field("strides", &b.strides())
86        .field("suboffsets", &b.suboffsets())
87        .field("internal", &raw.internal)
88        .finish()
89}
90
91/// Represents the type of a Python buffer element.
92#[derive(Copy, Clone, Debug, Eq, PartialEq)]
93pub enum ElementType {
94    /// A signed integer type.
95    SignedInteger {
96        /// The width of the signed integer in bytes.
97        bytes: usize,
98    },
99    /// An unsigned integer type.
100    UnsignedInteger {
101        /// The width of the unsigned integer in bytes.
102        bytes: usize,
103    },
104    /// A boolean type.
105    Bool,
106    /// A float type.
107    Float {
108        /// The width of the float in bytes.
109        bytes: usize,
110    },
111    /// An unknown type. This may occur when parsing has failed.
112    Unknown,
113}
114
115impl ElementType {
116    /// Determines the `ElementType` from a Python `struct` module format string.
117    ///
118    /// See <https://docs.python.org/3/library/struct.html#format-strings> for more information
119    /// about struct format strings.
120    pub fn from_format(format: &CStr) -> ElementType {
121        match format.to_bytes() {
122            [size] | [b'@', size] => native_element_type_from_type_char(*size),
123            [b'=' | b'<' | b'>' | b'!', size] => standard_element_type_from_type_char(*size),
124            _ => ElementType::Unknown,
125        }
126    }
127}
128
129fn native_element_type_from_type_char(type_char: u8) -> ElementType {
130    use self::ElementType::*;
131    match type_char {
132        b'c' => UnsignedInteger {
133            bytes: mem::size_of::<c_char>(),
134        },
135        b'b' => SignedInteger {
136            bytes: mem::size_of::<c_schar>(),
137        },
138        b'B' => UnsignedInteger {
139            bytes: mem::size_of::<c_uchar>(),
140        },
141        b'?' => Bool,
142        b'h' => SignedInteger {
143            bytes: mem::size_of::<c_short>(),
144        },
145        b'H' => UnsignedInteger {
146            bytes: mem::size_of::<c_ushort>(),
147        },
148        b'i' => SignedInteger {
149            bytes: mem::size_of::<c_int>(),
150        },
151        b'I' => UnsignedInteger {
152            bytes: mem::size_of::<c_uint>(),
153        },
154        b'l' => SignedInteger {
155            bytes: mem::size_of::<c_long>(),
156        },
157        b'L' => UnsignedInteger {
158            bytes: mem::size_of::<c_ulong>(),
159        },
160        b'q' => SignedInteger {
161            bytes: mem::size_of::<c_longlong>(),
162        },
163        b'Q' => UnsignedInteger {
164            bytes: mem::size_of::<c_ulonglong>(),
165        },
166        b'n' => SignedInteger {
167            bytes: mem::size_of::<libc::ssize_t>(),
168        },
169        b'N' => UnsignedInteger {
170            bytes: mem::size_of::<libc::size_t>(),
171        },
172        b'e' => Float { bytes: 2 },
173        b'f' => Float { bytes: 4 },
174        b'd' => Float { bytes: 8 },
175        _ => Unknown,
176    }
177}
178
179fn standard_element_type_from_type_char(type_char: u8) -> ElementType {
180    use self::ElementType::*;
181    match type_char {
182        b'c' | b'B' => UnsignedInteger { bytes: 1 },
183        b'b' => SignedInteger { bytes: 1 },
184        b'?' => Bool,
185        b'h' => SignedInteger { bytes: 2 },
186        b'H' => UnsignedInteger { bytes: 2 },
187        b'i' | b'l' => SignedInteger { bytes: 4 },
188        b'I' | b'L' => UnsignedInteger { bytes: 4 },
189        b'q' => SignedInteger { bytes: 8 },
190        b'Q' => UnsignedInteger { bytes: 8 },
191        b'e' => Float { bytes: 2 },
192        b'f' => Float { bytes: 4 },
193        b'd' => Float { bytes: 8 },
194        _ => Unknown,
195    }
196}
197
198#[cfg(target_endian = "little")]
199fn is_matching_endian(c: u8) -> bool {
200    c == b'@' || c == b'=' || c == b'>'
201}
202
203#[cfg(target_endian = "big")]
204fn is_matching_endian(c: u8) -> bool {
205    c == b'@' || c == b'=' || c == b'>' || c == b'!'
206}
207
208/// Trait implemented for possible element types of `PyBuffer`.
209///
210/// # Safety
211///
212/// This trait must only be implemented for types which represent valid elements of Python buffers.
213pub unsafe trait Element: Copy {
214    /// Gets whether the element specified in the format string is potentially compatible.
215    /// Alignment and size are checked separately from this function.
216    fn is_compatible_format(format: &CStr) -> bool;
217}
218
219impl<T: Element> FromPyObject<'_, '_> for PyBuffer<T> {
220    type Error = PyErr;
221
222    #[cfg(feature = "experimental-inspect")]
223    const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("collections.abc", "Buffer");
224
225    fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result<PyBuffer<T>, Self::Error> {
226        Self::get(&obj)
227    }
228}
229
230impl<T: Element> PyBuffer<T> {
231    /// Gets the underlying buffer from the specified python object.
232    pub fn get(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
233        PyUntypedBuffer::get(obj)?.into_typed()
234    }
235
236    /// Gets the buffer memory as a slice.
237    ///
238    /// This function succeeds if:
239    /// * the buffer format is compatible with `T`
240    /// * alignment and size of buffer elements is matching the expectations for type `T`
241    /// * the buffer is C-style contiguous
242    ///
243    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
244    /// to modify the values in the slice.
245    pub fn as_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [ReadOnlyCell<T>]> {
246        if self.is_c_contiguous() {
247            unsafe {
248                Some(slice::from_raw_parts(
249                    self.raw().buf.cast(),
250                    self.item_count(),
251                ))
252            }
253        } else {
254            None
255        }
256    }
257
258    /// Gets the buffer memory as a slice.
259    ///
260    /// This function succeeds if:
261    /// * the buffer is not read-only
262    /// * the buffer format is compatible with `T`
263    /// * alignment and size of buffer elements is matching the expectations for type `T`
264    /// * the buffer is C-style contiguous
265    ///
266    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
267    /// to modify the values in the slice.
268    pub fn as_mut_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [cell::Cell<T>]> {
269        if !self.readonly() && self.is_c_contiguous() {
270            unsafe {
271                Some(slice::from_raw_parts(
272                    self.raw().buf.cast(),
273                    self.item_count(),
274                ))
275            }
276        } else {
277            None
278        }
279    }
280
281    /// Gets the buffer memory as a slice.
282    ///
283    /// This function succeeds if:
284    /// * the buffer format is compatible with `T`
285    /// * alignment and size of buffer elements is matching the expectations for type `T`
286    /// * the buffer is Fortran-style contiguous
287    ///
288    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
289    /// to modify the values in the slice.
290    pub fn as_fortran_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [ReadOnlyCell<T>]> {
291        if mem::size_of::<T>() == self.item_size() && self.is_fortran_contiguous() {
292            unsafe {
293                Some(slice::from_raw_parts(
294                    self.raw().buf.cast(),
295                    self.item_count(),
296                ))
297            }
298        } else {
299            None
300        }
301    }
302
303    /// Gets the buffer memory as a slice.
304    ///
305    /// This function succeeds if:
306    /// * the buffer is not read-only
307    /// * the buffer format is compatible with `T`
308    /// * alignment and size of buffer elements is matching the expectations for type `T`
309    /// * the buffer is Fortran-style contiguous
310    ///
311    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
312    /// to modify the values in the slice.
313    pub fn as_fortran_mut_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [cell::Cell<T>]> {
314        if !self.readonly() && self.is_fortran_contiguous() {
315            unsafe {
316                Some(slice::from_raw_parts(
317                    self.raw().buf.cast(),
318                    self.item_count(),
319                ))
320            }
321        } else {
322            None
323        }
324    }
325
326    /// Copies the buffer elements to the specified slice.
327    /// If the buffer is multi-dimensional, the elements are written in C-style order.
328    ///
329    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
330    ///  * Fails if the buffer format is not compatible with type `T`.
331    ///
332    /// To check whether the buffer format is compatible before calling this method,
333    /// you can use `<T as buffer::Element>::is_compatible_format(buf.format())`.
334    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
335    pub fn copy_to_slice(&self, py: Python<'_>, target: &mut [T]) -> PyResult<()> {
336        self._copy_to_slice(py, target, b'C')
337    }
338
339    /// Copies the buffer elements to the specified slice.
340    /// If the buffer is multi-dimensional, the elements are written in Fortran-style order.
341    ///
342    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
343    ///  * Fails if the buffer format is not compatible with type `T`.
344    ///
345    /// To check whether the buffer format is compatible before calling this method,
346    /// you can use `<T as buffer::Element>::is_compatible_format(buf.format())`.
347    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
348    pub fn copy_to_fortran_slice(&self, py: Python<'_>, target: &mut [T]) -> PyResult<()> {
349        self._copy_to_slice(py, target, b'F')
350    }
351
352    fn _copy_to_slice(&self, py: Python<'_>, target: &mut [T], fort: u8) -> PyResult<()> {
353        if mem::size_of_val(target) != self.len_bytes() {
354            return Err(PyBufferError::new_err(format!(
355                "slice to copy to (of length {}) does not match buffer length of {}",
356                target.len(),
357                self.item_count()
358            )));
359        }
360
361        err::error_on_minusone(py, unsafe {
362            ffi::PyBuffer_ToContiguous(
363                target.as_mut_ptr().cast(),
364                #[cfg(Py_3_11)]
365                self.raw(),
366                #[cfg(not(Py_3_11))]
367                ptr::from_ref(self.raw()).cast_mut(),
368                self.raw().len,
369                fort as std::ffi::c_char,
370            )
371        })
372    }
373
374    /// Copies the buffer elements to a newly allocated vector.
375    /// If the buffer is multi-dimensional, the elements are written in C-style order.
376    ///
377    /// Fails if the buffer format is not compatible with type `T`.
378    pub fn to_vec(&self, py: Python<'_>) -> PyResult<Vec<T>> {
379        self._to_vec(py, b'C')
380    }
381
382    /// Copies the buffer elements to a newly allocated vector.
383    /// If the buffer is multi-dimensional, the elements are written in Fortran-style order.
384    ///
385    /// Fails if the buffer format is not compatible with type `T`.
386    pub fn to_fortran_vec(&self, py: Python<'_>) -> PyResult<Vec<T>> {
387        self._to_vec(py, b'F')
388    }
389
390    fn _to_vec(&self, py: Python<'_>, fort: u8) -> PyResult<Vec<T>> {
391        let item_count = self.item_count();
392        let mut vec: Vec<T> = Vec::with_capacity(item_count);
393
394        // Copy the buffer into the uninitialized space in the vector.
395        // Due to T:Copy, we don't need to be concerned with Drop impls.
396        err::error_on_minusone(py, unsafe {
397            ffi::PyBuffer_ToContiguous(
398                vec.as_mut_ptr().cast(),
399                #[cfg(Py_3_11)]
400                self.raw(),
401                #[cfg(not(Py_3_11))]
402                ptr::from_ref(self.raw()).cast_mut(),
403                self.raw().len,
404                fort as std::ffi::c_char,
405            )
406        })?;
407        // set vector length to mark the now-initialized space as usable
408        unsafe { vec.set_len(item_count) };
409        Ok(vec)
410    }
411
412    /// Copies the specified slice into the buffer.
413    /// If the buffer is multi-dimensional, the elements in the slice are expected to be in C-style order.
414    ///
415    ///  * Fails if the buffer is read-only.
416    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
417    ///  * Fails if the buffer format is not compatible with type `T`.
418    ///
419    /// To check whether the buffer format is compatible before calling this method,
420    /// use `<T as buffer::Element>::is_compatible_format(buf.format())`.
421    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
422    pub fn copy_from_slice(&self, py: Python<'_>, source: &[T]) -> PyResult<()> {
423        self._copy_from_slice(py, source, b'C')
424    }
425
426    /// Copies the specified slice into the buffer.
427    /// If the buffer is multi-dimensional, the elements in the slice are expected to be in Fortran-style order.
428    ///
429    ///  * Fails if the buffer is read-only.
430    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
431    ///  * Fails if the buffer format is not compatible with type `T`.
432    ///
433    /// To check whether the buffer format is compatible before calling this method,
434    /// use `<T as buffer::Element>::is_compatible_format(buf.format())`.
435    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
436    pub fn copy_from_fortran_slice(&self, py: Python<'_>, source: &[T]) -> PyResult<()> {
437        self._copy_from_slice(py, source, b'F')
438    }
439
440    fn _copy_from_slice(&self, py: Python<'_>, source: &[T], fort: u8) -> PyResult<()> {
441        if self.readonly() {
442            return Err(PyBufferError::new_err("cannot write to read-only buffer"));
443        } else if mem::size_of_val(source) != self.len_bytes() {
444            return Err(PyBufferError::new_err(format!(
445                "slice to copy from (of length {}) does not match buffer length of {}",
446                source.len(),
447                self.item_count()
448            )));
449        }
450
451        err::error_on_minusone(py, unsafe {
452            ffi::PyBuffer_FromContiguous(
453                #[cfg(Py_3_11)]
454                self.raw(),
455                #[cfg(not(Py_3_11))]
456                ptr::from_ref(self.raw()).cast_mut(),
457                #[cfg(Py_3_11)]
458                {
459                    source.as_ptr().cast()
460                },
461                #[cfg(not(Py_3_11))]
462                {
463                    source.as_ptr().cast::<c_void>().cast_mut()
464                },
465                self.raw().len,
466                fort as std::ffi::c_char,
467            )
468        })
469    }
470}
471
472impl<T> std::ops::Deref for PyBuffer<T> {
473    type Target = PyUntypedBuffer;
474
475    fn deref(&self) -> &Self::Target {
476        &self.0
477    }
478}
479
480impl PyUntypedBuffer {
481    /// Gets the underlying buffer from the specified python object.
482    pub fn get(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
483        let buf = {
484            let mut buf = Box::<RawBuffer>::new_uninit();
485            // SAFETY: RawBuffer is `#[repr(transparent)]` around FFI struct
486            err::error_on_minusone(obj.py(), unsafe {
487                ffi::PyObject_GetBuffer(
488                    obj.as_ptr(),
489                    buf.as_mut_ptr().cast::<ffi::Py_buffer>(),
490                    ffi::PyBUF_FULL_RO,
491                )
492            })?;
493            // Safety: buf is initialized by PyObject_GetBuffer.
494            unsafe { buf.assume_init() }
495        };
496        // Create PyBuffer immediately so that if validation checks fail, the PyBuffer::drop code
497        // will call PyBuffer_Release (thus avoiding any leaks).
498        let buf = Self(Pin::from(buf));
499        let raw = buf.raw();
500
501        if raw.shape.is_null() {
502            Err(PyBufferError::new_err("shape is null"))
503        } else if raw.strides.is_null() {
504            Err(PyBufferError::new_err("strides is null"))
505        } else {
506            Ok(buf)
507        }
508    }
509
510    /// Returns a `[PyBuffer]` instance if the buffer can be interpreted as containing elements of type `T`.
511    pub fn into_typed<T: Element>(self) -> PyResult<PyBuffer<T>> {
512        self.ensure_compatible_with::<T>()?;
513        Ok(PyBuffer(self, PhantomData))
514    }
515
516    /// Non-owning equivalent of [`into_typed()`][Self::into_typed].
517    pub fn as_typed<T: Element>(&self) -> PyResult<&PyBuffer<T>> {
518        self.ensure_compatible_with::<T>()?;
519        // SAFETY: PyBuffer<T> is repr(transparent) around PyUntypedBuffer
520        Ok(unsafe { NonNull::from(self).cast::<PyBuffer<T>>().as_ref() })
521    }
522
523    fn ensure_compatible_with<T: Element>(&self) -> PyResult<()> {
524        if mem::size_of::<T>() != self.item_size() || !T::is_compatible_format(self.format()) {
525            Err(PyBufferError::new_err(format!(
526                "buffer contents are not compatible with {}",
527                std::any::type_name::<T>()
528            )))
529        } else if self.raw().buf.align_offset(mem::align_of::<T>()) != 0 {
530            Err(PyBufferError::new_err(format!(
531                "buffer contents are insufficiently aligned for {}",
532                std::any::type_name::<T>()
533            )))
534        } else {
535            Ok(())
536        }
537    }
538
539    /// Releases the buffer object, freeing the reference to the Python object
540    /// which owns the buffer.
541    ///
542    /// This will automatically be called on drop.
543    pub fn release(self, _py: Python<'_>) {
544        // First move self into a ManuallyDrop, so that PyBuffer::drop will
545        // never be called. (It would attach to the interpreter and call PyBuffer_Release
546        // again.)
547        let mut mdself = mem::ManuallyDrop::new(self);
548        unsafe {
549            // Next, make the actual PyBuffer_Release call.
550            // Fine to get a mutable reference to the inner ffi::Py_buffer here, as we're destroying it.
551            mdself.0.release();
552
553            // Finally, drop the contained Pin<Box<_>> in place, to free the
554            // Box memory.
555            ptr::drop_in_place::<Pin<Box<RawBuffer>>>(&mut mdself.0);
556        }
557    }
558
559    /// Gets the pointer to the start of the buffer memory.
560    ///
561    /// Warning: the buffer memory can be mutated by other code (including
562    /// other Python functions, if the GIL is released, or other extension
563    /// modules even if the GIL is held). You must either access memory
564    /// atomically, or ensure there are no data races yourself. See
565    /// [this blog post] for more details.
566    ///
567    /// [this blog post]: https://alexgaynor.net/2022/oct/23/buffers-on-the-edge/
568    #[inline]
569    pub fn buf_ptr(&self) -> *mut c_void {
570        self.raw().buf
571    }
572
573    /// Gets a pointer to the specified item.
574    ///
575    /// If `indices.len() < self.dimensions()`, returns the start address of the sub-array at the specified dimension.
576    pub fn get_ptr(&self, indices: &[usize]) -> *mut c_void {
577        let shape = &self.shape()[..indices.len()];
578        for i in 0..indices.len() {
579            assert!(indices[i] < shape[i]);
580        }
581        unsafe {
582            ffi::PyBuffer_GetPointer(
583                #[cfg(Py_3_11)]
584                self.raw(),
585                #[cfg(not(Py_3_11))]
586                ptr::from_ref(self.raw()).cast_mut(),
587                #[cfg(Py_3_11)]
588                indices.as_ptr().cast(),
589                #[cfg(not(Py_3_11))]
590                indices.as_ptr().cast_mut().cast(),
591            )
592        }
593    }
594
595    /// Gets whether the underlying buffer is read-only.
596    #[inline]
597    pub fn readonly(&self) -> bool {
598        self.raw().readonly != 0
599    }
600
601    /// Gets the size of a single element, in bytes.
602    /// Important exception: when requesting an unformatted buffer, item_size still has the value
603    #[inline]
604    pub fn item_size(&self) -> usize {
605        self.raw().itemsize as usize
606    }
607
608    /// Gets the total number of items.
609    #[inline]
610    pub fn item_count(&self) -> usize {
611        (self.raw().len as usize) / (self.raw().itemsize as usize)
612    }
613
614    /// `item_size() * item_count()`.
615    /// For contiguous arrays, this is the length of the underlying memory block.
616    /// For non-contiguous arrays, it is the length that the logical structure would have if it were copied to a contiguous representation.
617    #[inline]
618    pub fn len_bytes(&self) -> usize {
619        self.raw().len as usize
620    }
621
622    /// Gets the number of dimensions.
623    ///
624    /// May be 0 to indicate a single scalar value.
625    #[inline]
626    pub fn dimensions(&self) -> usize {
627        self.raw().ndim as usize
628    }
629
630    /// Returns an array of length `dimensions`. `shape()[i]` is the length of the array in dimension number `i`.
631    ///
632    /// May return None for single-dimensional arrays or scalar values (`dimensions() <= 1`);
633    /// You can call `item_count()` to get the length of the single dimension.
634    ///
635    /// Despite Python using an array of signed integers, the values are guaranteed to be non-negative.
636    /// However, dimensions of length 0 are possible and might need special attention.
637    #[inline]
638    pub fn shape(&self) -> &[usize] {
639        unsafe { slice::from_raw_parts(self.raw().shape.cast(), self.raw().ndim as usize) }
640    }
641
642    /// Returns an array that holds, for each dimension, the number of bytes to skip to get to the next element in the dimension.
643    ///
644    /// Stride values can be any integer. For regular arrays, strides are usually positive,
645    /// but a consumer MUST be able to handle the case `strides[n] <= 0`.
646    #[inline]
647    pub fn strides(&self) -> &[isize] {
648        unsafe { slice::from_raw_parts(self.raw().strides, self.raw().ndim as usize) }
649    }
650
651    /// An array of length ndim.
652    /// If `suboffsets[n] >= 0`, the values stored along the nth dimension are pointers and the suboffset value dictates how many bytes to add to each pointer after de-referencing.
653    /// A suboffset value that is negative indicates that no de-referencing should occur (striding in a contiguous memory block).
654    ///
655    /// If all suboffsets are negative (i.e. no de-referencing is needed), then this field must be NULL (the default value).
656    #[inline]
657    pub fn suboffsets(&self) -> Option<&[isize]> {
658        unsafe {
659            if self.raw().suboffsets.is_null() {
660                None
661            } else {
662                Some(slice::from_raw_parts(
663                    self.raw().suboffsets,
664                    self.raw().ndim as usize,
665                ))
666            }
667        }
668    }
669
670    /// A string in struct module style syntax describing the contents of a single item.
671    #[inline]
672    pub fn format(&self) -> &CStr {
673        if self.raw().format.is_null() {
674            ffi::c_str!("B")
675        } else {
676            unsafe { CStr::from_ptr(self.raw().format) }
677        }
678    }
679
680    /// Gets whether the buffer is contiguous in C-style order (last index varies fastest when visiting items in order of memory address).
681    #[inline]
682    pub fn is_c_contiguous(&self) -> bool {
683        unsafe { ffi::PyBuffer_IsContiguous(self.raw(), b'C' as std::ffi::c_char) != 0 }
684    }
685
686    /// Gets whether the buffer is contiguous in Fortran-style order (first index varies fastest when visiting items in order of memory address).
687    #[inline]
688    pub fn is_fortran_contiguous(&self) -> bool {
689        unsafe { ffi::PyBuffer_IsContiguous(self.raw(), b'F' as std::ffi::c_char) != 0 }
690    }
691
692    fn raw(&self) -> &ffi::Py_buffer {
693        &self.0 .0
694    }
695}
696
697impl RawBuffer {
698    /// Release the contents of this pinned buffer.
699    ///
700    /// # Safety
701    ///
702    /// - The buffer must not be used after calling this function.
703    /// - This function can only be called once.
704    /// - Must be attached to the interpreter.
705    ///
706    unsafe fn release(self: &mut Pin<Box<Self>>) {
707        unsafe {
708            ffi::PyBuffer_Release(&mut Pin::get_unchecked_mut(self.as_mut()).0);
709        }
710    }
711}
712
713impl Drop for PyUntypedBuffer {
714    fn drop(&mut self) {
715        if Python::try_attach(|_| unsafe { self.0.release() }).is_none()
716            && crate::internal::state::is_in_gc_traversal()
717        {
718            eprintln!("Warning: PyBuffer dropped while in GC traversal, this is a bug and will leak memory.");
719        }
720        // If `try_attach` failed and `is_in_gc_traversal()` is false, then probably the interpreter has
721        // already finalized and we can just assume that the underlying memory has already been freed.
722        //
723        // So we don't handle that case here.
724    }
725}
726
727/// Like [std::cell::Cell], but only provides read-only access to the data.
728///
729/// `&ReadOnlyCell<T>` is basically a safe version of `*const T`:
730///  The data cannot be modified through the reference, but other references may
731///  be modifying the data.
732#[repr(transparent)]
733pub struct ReadOnlyCell<T: Element>(cell::UnsafeCell<T>);
734
735impl<T: Element> ReadOnlyCell<T> {
736    /// Returns a copy of the current value.
737    #[inline]
738    pub fn get(&self) -> T {
739        unsafe { *self.0.get() }
740    }
741
742    /// Returns a pointer to the current value.
743    #[inline]
744    pub fn as_ptr(&self) -> *const T {
745        self.0.get()
746    }
747}
748
749macro_rules! impl_element(
750    ($t:ty, $f:ident) => {
751        unsafe impl Element for $t {
752            fn is_compatible_format(format: &CStr) -> bool {
753                let slice = format.to_bytes();
754                if slice.len() > 1 && !is_matching_endian(slice[0]) {
755                    return false;
756                }
757                ElementType::from_format(format) == ElementType::$f { bytes: mem::size_of::<$t>() }
758            }
759        }
760    }
761);
762
763impl_element!(u8, UnsignedInteger);
764impl_element!(u16, UnsignedInteger);
765impl_element!(u32, UnsignedInteger);
766impl_element!(u64, UnsignedInteger);
767impl_element!(usize, UnsignedInteger);
768impl_element!(i8, SignedInteger);
769impl_element!(i16, SignedInteger);
770impl_element!(i32, SignedInteger);
771impl_element!(i64, SignedInteger);
772impl_element!(isize, SignedInteger);
773impl_element!(f32, Float);
774impl_element!(f64, Float);
775
776#[cfg(test)]
777mod tests {
778    use super::*;
779
780    use crate::ffi;
781    use crate::types::any::PyAnyMethods;
782    use crate::types::PyBytes;
783    use crate::Python;
784
785    #[test]
786    fn test_debug() {
787        Python::attach(|py| {
788            let bytes = PyBytes::new(py, b"abcde");
789            let buffer: PyBuffer<u8> = PyBuffer::get(&bytes).unwrap();
790            let expected = format!(
791                concat!(
792                    "PyBuffer {{ buf: {:?}, obj: {:?}, ",
793                    "len: 5, itemsize: 1, readonly: 1, ",
794                    "ndim: 1, format: \"B\", shape: [5], ",
795                    "strides: [1], suboffsets: None, internal: {:?} }}",
796                ),
797                buffer.raw().buf,
798                buffer.raw().obj,
799                buffer.raw().internal
800            );
801            let debug_repr = format!("{:?}", buffer);
802            assert_eq!(debug_repr, expected);
803        });
804    }
805
806    #[test]
807    fn test_element_type_from_format() {
808        use super::ElementType::*;
809        use std::mem::size_of;
810
811        for (cstr, expected) in [
812            // @ prefix goes to native_element_type_from_type_char
813            (
814                c"@b",
815                SignedInteger {
816                    bytes: size_of::<c_schar>(),
817                },
818            ),
819            (
820                c"@c",
821                UnsignedInteger {
822                    bytes: size_of::<c_char>(),
823                },
824            ),
825            (
826                c"@b",
827                SignedInteger {
828                    bytes: size_of::<c_schar>(),
829                },
830            ),
831            (
832                c"@B",
833                UnsignedInteger {
834                    bytes: size_of::<c_uchar>(),
835                },
836            ),
837            (c"@?", Bool),
838            (
839                c"@h",
840                SignedInteger {
841                    bytes: size_of::<c_short>(),
842                },
843            ),
844            (
845                c"@H",
846                UnsignedInteger {
847                    bytes: size_of::<c_ushort>(),
848                },
849            ),
850            (
851                c"@i",
852                SignedInteger {
853                    bytes: size_of::<c_int>(),
854                },
855            ),
856            (
857                c"@I",
858                UnsignedInteger {
859                    bytes: size_of::<c_uint>(),
860                },
861            ),
862            (
863                c"@l",
864                SignedInteger {
865                    bytes: size_of::<c_long>(),
866                },
867            ),
868            (
869                c"@L",
870                UnsignedInteger {
871                    bytes: size_of::<c_ulong>(),
872                },
873            ),
874            (
875                c"@q",
876                SignedInteger {
877                    bytes: size_of::<c_longlong>(),
878                },
879            ),
880            (
881                c"@Q",
882                UnsignedInteger {
883                    bytes: size_of::<c_ulonglong>(),
884                },
885            ),
886            (
887                c"@n",
888                SignedInteger {
889                    bytes: size_of::<libc::ssize_t>(),
890                },
891            ),
892            (
893                c"@N",
894                UnsignedInteger {
895                    bytes: size_of::<libc::size_t>(),
896                },
897            ),
898            (c"@e", Float { bytes: 2 }),
899            (c"@f", Float { bytes: 4 }),
900            (c"@d", Float { bytes: 8 }),
901            (c"@z", Unknown),
902            // = prefix goes to standard_element_type_from_type_char
903            (c"=b", SignedInteger { bytes: 1 }),
904            (c"=c", UnsignedInteger { bytes: 1 }),
905            (c"=B", UnsignedInteger { bytes: 1 }),
906            (c"=?", Bool),
907            (c"=h", SignedInteger { bytes: 2 }),
908            (c"=H", UnsignedInteger { bytes: 2 }),
909            (c"=l", SignedInteger { bytes: 4 }),
910            (c"=l", SignedInteger { bytes: 4 }),
911            (c"=I", UnsignedInteger { bytes: 4 }),
912            (c"=L", UnsignedInteger { bytes: 4 }),
913            (c"=q", SignedInteger { bytes: 8 }),
914            (c"=Q", UnsignedInteger { bytes: 8 }),
915            (c"=e", Float { bytes: 2 }),
916            (c"=f", Float { bytes: 4 }),
917            (c"=d", Float { bytes: 8 }),
918            (c"=z", Unknown),
919            (c"=0", Unknown),
920            // unknown prefix -> Unknown
921            (c":b", Unknown),
922        ] {
923            assert_eq!(
924                ElementType::from_format(cstr),
925                expected,
926                "element from format &Cstr: {cstr:?}",
927            );
928        }
929    }
930
931    #[test]
932    fn test_compatible_size() {
933        // for the cast in PyBuffer::shape()
934        assert_eq!(
935            std::mem::size_of::<ffi::Py_ssize_t>(),
936            std::mem::size_of::<usize>()
937        );
938    }
939
940    #[test]
941    fn test_bytes_buffer() {
942        Python::attach(|py| {
943            let bytes = PyBytes::new(py, b"abcde");
944            let buffer = PyBuffer::get(&bytes).unwrap();
945            assert_eq!(buffer.dimensions(), 1);
946            assert_eq!(buffer.item_count(), 5);
947            assert_eq!(buffer.format().to_str().unwrap(), "B");
948            assert_eq!(buffer.shape(), [5]);
949            // single-dimensional buffer is always contiguous
950            assert!(buffer.is_c_contiguous());
951            assert!(buffer.is_fortran_contiguous());
952
953            let slice = buffer.as_slice(py).unwrap();
954            assert_eq!(slice.len(), 5);
955            assert_eq!(slice[0].get(), b'a');
956            assert_eq!(slice[2].get(), b'c');
957
958            assert_eq!(unsafe { *(buffer.get_ptr(&[1]).cast::<u8>()) }, b'b');
959
960            assert!(buffer.as_mut_slice(py).is_none());
961
962            assert!(buffer.copy_to_slice(py, &mut [0u8]).is_err());
963            let mut arr = [0; 5];
964            buffer.copy_to_slice(py, &mut arr).unwrap();
965            assert_eq!(arr, b"abcde" as &[u8]);
966
967            assert!(buffer.copy_from_slice(py, &[0u8; 5]).is_err());
968            assert_eq!(buffer.to_vec(py).unwrap(), b"abcde");
969        });
970    }
971
972    #[test]
973    fn test_array_buffer() {
974        Python::attach(|py| {
975            let array = py
976                .import("array")
977                .unwrap()
978                .call_method("array", ("f", (1.0, 1.5, 2.0, 2.5)), None)
979                .unwrap();
980            let buffer = PyBuffer::get(&array).unwrap();
981            assert_eq!(buffer.dimensions(), 1);
982            assert_eq!(buffer.item_count(), 4);
983            assert_eq!(buffer.format().to_str().unwrap(), "f");
984            assert_eq!(buffer.shape(), [4]);
985
986            // array creates a 1D contiguous buffer, so it's both C and F contiguous.  This would
987            // be more interesting if we can come up with a 2D buffer but I think it would need a
988            // third-party lib or a custom class.
989
990            // C-contiguous fns
991            let slice = buffer.as_slice(py).unwrap();
992            assert_eq!(slice.len(), 4);
993            assert_eq!(slice[0].get(), 1.0);
994            assert_eq!(slice[3].get(), 2.5);
995
996            let mut_slice = buffer.as_mut_slice(py).unwrap();
997            assert_eq!(mut_slice.len(), 4);
998            assert_eq!(mut_slice[0].get(), 1.0);
999            mut_slice[3].set(2.75);
1000            assert_eq!(slice[3].get(), 2.75);
1001
1002            buffer
1003                .copy_from_slice(py, &[10.0f32, 11.0, 12.0, 13.0])
1004                .unwrap();
1005            assert_eq!(slice[2].get(), 12.0);
1006
1007            assert_eq!(buffer.to_vec(py).unwrap(), [10.0, 11.0, 12.0, 13.0]);
1008
1009            // F-contiguous fns
1010            let buffer = PyBuffer::get(&array).unwrap();
1011            let slice = buffer.as_fortran_slice(py).unwrap();
1012            assert_eq!(slice.len(), 4);
1013            assert_eq!(slice[1].get(), 11.0);
1014
1015            let mut_slice = buffer.as_fortran_mut_slice(py).unwrap();
1016            assert_eq!(mut_slice.len(), 4);
1017            assert_eq!(mut_slice[2].get(), 12.0);
1018            mut_slice[3].set(2.75);
1019            assert_eq!(slice[3].get(), 2.75);
1020
1021            buffer
1022                .copy_from_fortran_slice(py, &[10.0f32, 11.0, 12.0, 13.0])
1023                .unwrap();
1024            assert_eq!(slice[2].get(), 12.0);
1025
1026            assert_eq!(buffer.to_fortran_vec(py).unwrap(), [10.0, 11.0, 12.0, 13.0]);
1027        });
1028    }
1029
1030    #[test]
1031    fn test_untyped_buffer() {
1032        Python::attach(|py| {
1033            let bytes = PyBytes::new(py, b"abcde");
1034            let untyped = PyUntypedBuffer::get(&bytes).unwrap();
1035            assert_eq!(untyped.dimensions(), 1);
1036            assert_eq!(untyped.item_count(), 5);
1037            assert_eq!(untyped.format().to_str().unwrap(), "B");
1038            assert_eq!(untyped.shape(), [5]);
1039
1040            let typed: &PyBuffer<u8> = untyped.as_typed().unwrap();
1041            assert_eq!(typed.dimensions(), 1);
1042            assert_eq!(typed.item_count(), 5);
1043            assert_eq!(typed.format().to_str().unwrap(), "B");
1044            assert_eq!(typed.shape(), [5]);
1045        });
1046    }
1047}
pyo3/buffer.rs

pyo3/
buffer.rs
