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
21use crate::Bound;
22use crate::{err, exceptions::PyBufferError, ffi, FromPyObject, PyAny, PyResult, Python};
23use std::marker::PhantomData;
24use std::os::raw;
25use std::pin::Pin;
26use std::{cell, mem, ptr, slice};
27use std::{ffi::CStr, fmt::Debug};
28
29/// Allows access to the underlying buffer used by a python object such as `bytes`, `bytearray` or `array.array`.
30// use Pin<Box> because Python expects that the Py_buffer struct has a stable memory address
31#[repr(transparent)]
32pub struct PyBuffer<T>(Pin<Box<ffi::Py_buffer>>, PhantomData<T>);
33
34// PyBuffer is thread-safe: the shape of the buffer is immutable while a Py_buffer exists.
35// Accessing the buffer contents is protected using the GIL.
36unsafe impl<T> Send for PyBuffer<T> {}
37unsafe impl<T> Sync for PyBuffer<T> {}
38
39impl<T> Debug for PyBuffer<T> {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("PyBuffer")
42            .field("buf", &self.0.buf)
43            .field("obj", &self.0.obj)
44            .field("len", &self.0.len)
45            .field("itemsize", &self.0.itemsize)
46            .field("readonly", &self.0.readonly)
47            .field("ndim", &self.0.ndim)
48            .field("format", &self.0.format)
49            .field("shape", &self.0.shape)
50            .field("strides", &self.0.strides)
51            .field("suboffsets", &self.0.suboffsets)
52            .field("internal", &self.0.internal)
53            .finish()
54    }
55}
56
57/// Represents the type of a Python buffer element.
58#[derive(Copy, Clone, Debug, Eq, PartialEq)]
59pub enum ElementType {
60    /// A signed integer type.
61    SignedInteger {
62        /// The width of the signed integer in bytes.
63        bytes: usize,
64    },
65    /// An unsigned integer type.
66    UnsignedInteger {
67        /// The width of the unsigned integer in bytes.
68        bytes: usize,
69    },
70    /// A boolean type.
71    Bool,
72    /// A float type.
73    Float {
74        /// The width of the float in bytes.
75        bytes: usize,
76    },
77    /// An unknown type. This may occur when parsing has failed.
78    Unknown,
79}
80
81impl ElementType {
82    /// Determines the `ElementType` from a Python `struct` module format string.
83    ///
84    /// See <https://docs.python.org/3/library/struct.html#format-strings> for more information
85    /// about struct format strings.
86    pub fn from_format(format: &CStr) -> ElementType {
87        match format.to_bytes() {
88            [size] | [b'@', size] => native_element_type_from_type_char(*size),
89            [b'=' | b'<' | b'>' | b'!', size] => standard_element_type_from_type_char(*size),
90            _ => ElementType::Unknown,
91        }
92    }
93}
94
95fn native_element_type_from_type_char(type_char: u8) -> ElementType {
96    use self::ElementType::*;
97    match type_char {
98        b'c' => UnsignedInteger {
99            bytes: mem::size_of::<raw::c_char>(),
100        },
101        b'b' => SignedInteger {
102            bytes: mem::size_of::<raw::c_schar>(),
103        },
104        b'B' => UnsignedInteger {
105            bytes: mem::size_of::<raw::c_uchar>(),
106        },
107        b'?' => Bool,
108        b'h' => SignedInteger {
109            bytes: mem::size_of::<raw::c_short>(),
110        },
111        b'H' => UnsignedInteger {
112            bytes: mem::size_of::<raw::c_ushort>(),
113        },
114        b'i' => SignedInteger {
115            bytes: mem::size_of::<raw::c_int>(),
116        },
117        b'I' => UnsignedInteger {
118            bytes: mem::size_of::<raw::c_uint>(),
119        },
120        b'l' => SignedInteger {
121            bytes: mem::size_of::<raw::c_long>(),
122        },
123        b'L' => UnsignedInteger {
124            bytes: mem::size_of::<raw::c_ulong>(),
125        },
126        b'q' => SignedInteger {
127            bytes: mem::size_of::<raw::c_longlong>(),
128        },
129        b'Q' => UnsignedInteger {
130            bytes: mem::size_of::<raw::c_ulonglong>(),
131        },
132        b'n' => SignedInteger {
133            bytes: mem::size_of::<libc::ssize_t>(),
134        },
135        b'N' => UnsignedInteger {
136            bytes: mem::size_of::<libc::size_t>(),
137        },
138        b'e' => Float { bytes: 2 },
139        b'f' => Float { bytes: 4 },
140        b'd' => Float { bytes: 8 },
141        _ => Unknown,
142    }
143}
144
145fn standard_element_type_from_type_char(type_char: u8) -> ElementType {
146    use self::ElementType::*;
147    match type_char {
148        b'c' | b'B' => UnsignedInteger { bytes: 1 },
149        b'b' => SignedInteger { bytes: 1 },
150        b'?' => Bool,
151        b'h' => SignedInteger { bytes: 2 },
152        b'H' => UnsignedInteger { bytes: 2 },
153        b'i' | b'l' => SignedInteger { bytes: 4 },
154        b'I' | b'L' => UnsignedInteger { bytes: 4 },
155        b'q' => SignedInteger { bytes: 8 },
156        b'Q' => UnsignedInteger { bytes: 8 },
157        b'e' => Float { bytes: 2 },
158        b'f' => Float { bytes: 4 },
159        b'd' => Float { bytes: 8 },
160        _ => Unknown,
161    }
162}
163
164#[cfg(target_endian = "little")]
165fn is_matching_endian(c: u8) -> bool {
166    c == b'@' || c == b'=' || c == b'>'
167}
168
169#[cfg(target_endian = "big")]
170fn is_matching_endian(c: u8) -> bool {
171    c == b'@' || c == b'=' || c == b'>' || c == b'!'
172}
173
174/// Trait implemented for possible element types of `PyBuffer`.
175///
176/// # Safety
177///
178/// This trait must only be implemented for types which represent valid elements of Python buffers.
179pub unsafe trait Element: Copy {
180    /// Gets whether the element specified in the format string is potentially compatible.
181    /// Alignment and size are checked separately from this function.
182    fn is_compatible_format(format: &CStr) -> bool;
183}
184
185impl<T: Element> FromPyObject<'_> for PyBuffer<T> {
186    fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<PyBuffer<T>> {
187        Self::get(obj)
188    }
189}
190
191impl<T: Element> PyBuffer<T> {
192    /// Gets the underlying buffer from the specified python object.
193    pub fn get(obj: &Bound<'_, PyAny>) -> PyResult<PyBuffer<T>> {
194        // TODO: use nightly API Box::new_uninit() once our MSRV is 1.82
195        let mut buf = Box::new(mem::MaybeUninit::uninit());
196        let buf: Box<ffi::Py_buffer> = {
197            err::error_on_minusone(obj.py(), unsafe {
198                ffi::PyObject_GetBuffer(obj.as_ptr(), buf.as_mut_ptr(), ffi::PyBUF_FULL_RO)
199            })?;
200            // Safety: buf is initialized by PyObject_GetBuffer.
201            // TODO: use nightly API Box::assume_init() once our MSRV is 1.82
202            unsafe { mem::transmute(buf) }
203        };
204        // Create PyBuffer immediately so that if validation checks fail, the PyBuffer::drop code
205        // will call PyBuffer_Release (thus avoiding any leaks).
206        let buf = PyBuffer(Pin::from(buf), PhantomData);
207
208        if buf.0.shape.is_null() {
209            Err(PyBufferError::new_err("shape is null"))
210        } else if buf.0.strides.is_null() {
211            Err(PyBufferError::new_err("strides is null"))
212        } else if mem::size_of::<T>() != buf.item_size() || !T::is_compatible_format(buf.format()) {
213            Err(PyBufferError::new_err(format!(
214                "buffer contents are not compatible with {}",
215                std::any::type_name::<T>()
216            )))
217        } else if buf.0.buf.align_offset(mem::align_of::<T>()) != 0 {
218            Err(PyBufferError::new_err(format!(
219                "buffer contents are insufficiently aligned for {}",
220                std::any::type_name::<T>()
221            )))
222        } else {
223            Ok(buf)
224        }
225    }
226
227    /// Deprecated name for [`PyBuffer::get`].
228    #[deprecated(since = "0.23.0", note = "renamed to `PyBuffer::get`")]
229    #[inline]
230    pub fn get_bound(obj: &Bound<'_, PyAny>) -> PyResult<PyBuffer<T>> {
231        Self::get(obj)
232    }
233
234    /// Gets the pointer to the start of the buffer memory.
235    ///
236    /// Warning: the buffer memory might be mutated by other Python functions,
237    /// and thus may only be accessed while the GIL is held.
238    #[inline]
239    pub fn buf_ptr(&self) -> *mut raw::c_void {
240        self.0.buf
241    }
242
243    /// Gets a pointer to the specified item.
244    ///
245    /// If `indices.len() < self.dimensions()`, returns the start address of the sub-array at the specified dimension.
246    pub fn get_ptr(&self, indices: &[usize]) -> *mut raw::c_void {
247        let shape = &self.shape()[..indices.len()];
248        for i in 0..indices.len() {
249            assert!(indices[i] < shape[i]);
250        }
251        unsafe {
252            ffi::PyBuffer_GetPointer(
253                #[cfg(Py_3_11)]
254                &*self.0,
255                #[cfg(not(Py_3_11))]
256                {
257                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
258                },
259                #[cfg(Py_3_11)]
260                {
261                    indices.as_ptr().cast()
262                },
263                #[cfg(not(Py_3_11))]
264                {
265                    indices.as_ptr() as *mut ffi::Py_ssize_t
266                },
267            )
268        }
269    }
270
271    /// Gets whether the underlying buffer is read-only.
272    #[inline]
273    pub fn readonly(&self) -> bool {
274        self.0.readonly != 0
275    }
276
277    /// Gets the size of a single element, in bytes.
278    /// Important exception: when requesting an unformatted buffer, item_size still has the value
279    #[inline]
280    pub fn item_size(&self) -> usize {
281        self.0.itemsize as usize
282    }
283
284    /// Gets the total number of items.
285    #[inline]
286    pub fn item_count(&self) -> usize {
287        (self.0.len as usize) / (self.0.itemsize as usize)
288    }
289
290    /// `item_size() * item_count()`.
291    /// For contiguous arrays, this is the length of the underlying memory block.
292    /// For non-contiguous arrays, it is the length that the logical structure would have if it were copied to a contiguous representation.
293    #[inline]
294    pub fn len_bytes(&self) -> usize {
295        self.0.len as usize
296    }
297
298    /// Gets the number of dimensions.
299    ///
300    /// May be 0 to indicate a single scalar value.
301    #[inline]
302    pub fn dimensions(&self) -> usize {
303        self.0.ndim as usize
304    }
305
306    /// Returns an array of length `dimensions`. `shape()[i]` is the length of the array in dimension number `i`.
307    ///
308    /// May return None for single-dimensional arrays or scalar values (`dimensions() <= 1`);
309    /// You can call `item_count()` to get the length of the single dimension.
310    ///
311    /// Despite Python using an array of signed integers, the values are guaranteed to be non-negative.
312    /// However, dimensions of length 0 are possible and might need special attention.
313    #[inline]
314    pub fn shape(&self) -> &[usize] {
315        unsafe { slice::from_raw_parts(self.0.shape.cast(), self.0.ndim as usize) }
316    }
317
318    /// Returns an array that holds, for each dimension, the number of bytes to skip to get to the next element in the dimension.
319    ///
320    /// Stride values can be any integer. For regular arrays, strides are usually positive,
321    /// but a consumer MUST be able to handle the case `strides[n] <= 0`.
322    #[inline]
323    pub fn strides(&self) -> &[isize] {
324        unsafe { slice::from_raw_parts(self.0.strides, self.0.ndim as usize) }
325    }
326
327    /// An array of length ndim.
328    /// 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.
329    /// A suboffset value that is negative indicates that no de-referencing should occur (striding in a contiguous memory block).
330    ///
331    /// If all suboffsets are negative (i.e. no de-referencing is needed), then this field must be NULL (the default value).
332    #[inline]
333    pub fn suboffsets(&self) -> Option<&[isize]> {
334        unsafe {
335            if self.0.suboffsets.is_null() {
336                None
337            } else {
338                Some(slice::from_raw_parts(
339                    self.0.suboffsets,
340                    self.0.ndim as usize,
341                ))
342            }
343        }
344    }
345
346    /// A NUL terminated string in struct module style syntax describing the contents of a single item.
347    #[inline]
348    pub fn format(&self) -> &CStr {
349        if self.0.format.is_null() {
350            ffi::c_str!("B")
351        } else {
352            unsafe { CStr::from_ptr(self.0.format) }
353        }
354    }
355
356    /// Gets whether the buffer is contiguous in C-style order (last index varies fastest when visiting items in order of memory address).
357    #[inline]
358    pub fn is_c_contiguous(&self) -> bool {
359        unsafe { ffi::PyBuffer_IsContiguous(&*self.0, b'C' as std::os::raw::c_char) != 0 }
360    }
361
362    /// Gets whether the buffer is contiguous in Fortran-style order (first index varies fastest when visiting items in order of memory address).
363    #[inline]
364    pub fn is_fortran_contiguous(&self) -> bool {
365        unsafe { ffi::PyBuffer_IsContiguous(&*self.0, b'F' as std::os::raw::c_char) != 0 }
366    }
367
368    /// Gets the buffer memory as a slice.
369    ///
370    /// This function succeeds if:
371    /// * the buffer format is compatible with `T`
372    /// * alignment and size of buffer elements is matching the expectations for type `T`
373    /// * the buffer is C-style contiguous
374    ///
375    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
376    /// to modify the values in the slice.
377    pub fn as_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [ReadOnlyCell<T>]> {
378        if self.is_c_contiguous() {
379            unsafe {
380                Some(slice::from_raw_parts(
381                    self.0.buf as *mut ReadOnlyCell<T>,
382                    self.item_count(),
383                ))
384            }
385        } else {
386            None
387        }
388    }
389
390    /// Gets the buffer memory as a slice.
391    ///
392    /// This function succeeds if:
393    /// * the buffer is not read-only
394    /// * the buffer format is compatible with `T`
395    /// * alignment and size of buffer elements is matching the expectations for type `T`
396    /// * the buffer is C-style contiguous
397    ///
398    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
399    /// to modify the values in the slice.
400    pub fn as_mut_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [cell::Cell<T>]> {
401        if !self.readonly() && self.is_c_contiguous() {
402            unsafe {
403                Some(slice::from_raw_parts(
404                    self.0.buf as *mut cell::Cell<T>,
405                    self.item_count(),
406                ))
407            }
408        } else {
409            None
410        }
411    }
412
413    /// Gets the buffer memory as a slice.
414    ///
415    /// This function succeeds if:
416    /// * the buffer format is compatible with `T`
417    /// * alignment and size of buffer elements is matching the expectations for type `T`
418    /// * the buffer is Fortran-style contiguous
419    ///
420    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
421    /// to modify the values in the slice.
422    pub fn as_fortran_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [ReadOnlyCell<T>]> {
423        if mem::size_of::<T>() == self.item_size() && self.is_fortran_contiguous() {
424            unsafe {
425                Some(slice::from_raw_parts(
426                    self.0.buf as *mut ReadOnlyCell<T>,
427                    self.item_count(),
428                ))
429            }
430        } else {
431            None
432        }
433    }
434
435    /// Gets the buffer memory as a slice.
436    ///
437    /// This function succeeds if:
438    /// * the buffer is not read-only
439    /// * the buffer format is compatible with `T`
440    /// * alignment and size of buffer elements is matching the expectations for type `T`
441    /// * the buffer is Fortran-style contiguous
442    ///
443    /// The returned slice uses type `Cell<T>` because it's theoretically possible for any call into the Python runtime
444    /// to modify the values in the slice.
445    pub fn as_fortran_mut_slice<'a>(&'a self, _py: Python<'a>) -> Option<&'a [cell::Cell<T>]> {
446        if !self.readonly() && self.is_fortran_contiguous() {
447            unsafe {
448                Some(slice::from_raw_parts(
449                    self.0.buf as *mut cell::Cell<T>,
450                    self.item_count(),
451                ))
452            }
453        } else {
454            None
455        }
456    }
457
458    /// Copies the buffer elements to the specified slice.
459    /// If the buffer is multi-dimensional, the elements are written in C-style order.
460    ///
461    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
462    ///  * Fails if the buffer format is not compatible with type `T`.
463    ///
464    /// To check whether the buffer format is compatible before calling this method,
465    /// you can use `<T as buffer::Element>::is_compatible_format(buf.format())`.
466    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
467    pub fn copy_to_slice(&self, py: Python<'_>, target: &mut [T]) -> PyResult<()> {
468        self._copy_to_slice(py, target, b'C')
469    }
470
471    /// Copies the buffer elements to the specified slice.
472    /// If the buffer is multi-dimensional, the elements are written in Fortran-style order.
473    ///
474    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
475    ///  * Fails if the buffer format is not compatible with type `T`.
476    ///
477    /// To check whether the buffer format is compatible before calling this method,
478    /// you can use `<T as buffer::Element>::is_compatible_format(buf.format())`.
479    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
480    pub fn copy_to_fortran_slice(&self, py: Python<'_>, target: &mut [T]) -> PyResult<()> {
481        self._copy_to_slice(py, target, b'F')
482    }
483
484    fn _copy_to_slice(&self, py: Python<'_>, target: &mut [T], fort: u8) -> PyResult<()> {
485        if mem::size_of_val(target) != self.len_bytes() {
486            return Err(PyBufferError::new_err(format!(
487                "slice to copy to (of length {}) does not match buffer length of {}",
488                target.len(),
489                self.item_count()
490            )));
491        }
492
493        err::error_on_minusone(py, unsafe {
494            ffi::PyBuffer_ToContiguous(
495                target.as_mut_ptr().cast(),
496                #[cfg(Py_3_11)]
497                &*self.0,
498                #[cfg(not(Py_3_11))]
499                {
500                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
501                },
502                self.0.len,
503                fort as std::os::raw::c_char,
504            )
505        })
506    }
507
508    /// Copies the buffer elements to a newly allocated vector.
509    /// If the buffer is multi-dimensional, the elements are written in C-style order.
510    ///
511    /// Fails if the buffer format is not compatible with type `T`.
512    pub fn to_vec(&self, py: Python<'_>) -> PyResult<Vec<T>> {
513        self._to_vec(py, b'C')
514    }
515
516    /// Copies the buffer elements to a newly allocated vector.
517    /// If the buffer is multi-dimensional, the elements are written in Fortran-style order.
518    ///
519    /// Fails if the buffer format is not compatible with type `T`.
520    pub fn to_fortran_vec(&self, py: Python<'_>) -> PyResult<Vec<T>> {
521        self._to_vec(py, b'F')
522    }
523
524    fn _to_vec(&self, py: Python<'_>, fort: u8) -> PyResult<Vec<T>> {
525        let item_count = self.item_count();
526        let mut vec: Vec<T> = Vec::with_capacity(item_count);
527
528        // Copy the buffer into the uninitialized space in the vector.
529        // Due to T:Copy, we don't need to be concerned with Drop impls.
530        err::error_on_minusone(py, unsafe {
531            ffi::PyBuffer_ToContiguous(
532                vec.as_ptr() as *mut raw::c_void,
533                #[cfg(Py_3_11)]
534                &*self.0,
535                #[cfg(not(Py_3_11))]
536                {
537                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
538                },
539                self.0.len,
540                fort as std::os::raw::c_char,
541            )
542        })?;
543        // set vector length to mark the now-initialized space as usable
544        unsafe { vec.set_len(item_count) };
545        Ok(vec)
546    }
547
548    /// Copies the specified slice into the buffer.
549    /// If the buffer is multi-dimensional, the elements in the slice are expected to be in C-style order.
550    ///
551    ///  * Fails if the buffer is read-only.
552    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
553    ///  * Fails if the buffer format is not compatible with type `T`.
554    ///
555    /// To check whether the buffer format is compatible before calling this method,
556    /// use `<T as buffer::Element>::is_compatible_format(buf.format())`.
557    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
558    pub fn copy_from_slice(&self, py: Python<'_>, source: &[T]) -> PyResult<()> {
559        self._copy_from_slice(py, source, b'C')
560    }
561
562    /// Copies the specified slice into the buffer.
563    /// If the buffer is multi-dimensional, the elements in the slice are expected to be in Fortran-style order.
564    ///
565    ///  * Fails if the buffer is read-only.
566    ///  * Fails if the slice does not have the correct length (`buf.item_count()`).
567    ///  * Fails if the buffer format is not compatible with type `T`.
568    ///
569    /// To check whether the buffer format is compatible before calling this method,
570    /// use `<T as buffer::Element>::is_compatible_format(buf.format())`.
571    /// Alternatively, `match buffer::ElementType::from_format(buf.format())`.
572    pub fn copy_from_fortran_slice(&self, py: Python<'_>, source: &[T]) -> PyResult<()> {
573        self._copy_from_slice(py, source, b'F')
574    }
575
576    fn _copy_from_slice(&self, py: Python<'_>, source: &[T], fort: u8) -> PyResult<()> {
577        if self.readonly() {
578            return Err(PyBufferError::new_err("cannot write to read-only buffer"));
579        } else if mem::size_of_val(source) != self.len_bytes() {
580            return Err(PyBufferError::new_err(format!(
581                "slice to copy from (of length {}) does not match buffer length of {}",
582                source.len(),
583                self.item_count()
584            )));
585        }
586
587        err::error_on_minusone(py, unsafe {
588            ffi::PyBuffer_FromContiguous(
589                #[cfg(Py_3_11)]
590                &*self.0,
591                #[cfg(not(Py_3_11))]
592                {
593                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
594                },
595                #[cfg(Py_3_11)]
596                {
597                    source.as_ptr().cast()
598                },
599                #[cfg(not(Py_3_11))]
600                {
601                    source.as_ptr() as *mut raw::c_void
602                },
603                self.0.len,
604                fort as std::os::raw::c_char,
605            )
606        })
607    }
608
609    /// Releases the buffer object, freeing the reference to the Python object
610    /// which owns the buffer.
611    ///
612    /// This will automatically be called on drop.
613    pub fn release(self, _py: Python<'_>) {
614        // First move self into a ManuallyDrop, so that PyBuffer::drop will
615        // never be called. (It would acquire the GIL and call PyBuffer_Release
616        // again.)
617        let mut mdself = mem::ManuallyDrop::new(self);
618        unsafe {
619            // Next, make the actual PyBuffer_Release call.
620            ffi::PyBuffer_Release(&mut *mdself.0);
621
622            // Finally, drop the contained Pin<Box<_>> in place, to free the
623            // Box memory.
624            let inner: *mut Pin<Box<ffi::Py_buffer>> = &mut mdself.0;
625            ptr::drop_in_place(inner);
626        }
627    }
628}
629
630impl<T> Drop for PyBuffer<T> {
631    fn drop(&mut self) {
632        Python::with_gil(|_| unsafe { ffi::PyBuffer_Release(&mut *self.0) });
633    }
634}
635
636/// Like [std::cell::Cell], but only provides read-only access to the data.
637///
638/// `&ReadOnlyCell<T>` is basically a safe version of `*const T`:
639///  The data cannot be modified through the reference, but other references may
640///  be modifying the data.
641#[repr(transparent)]
642pub struct ReadOnlyCell<T: Element>(cell::UnsafeCell<T>);
643
644impl<T: Element> ReadOnlyCell<T> {
645    /// Returns a copy of the current value.
646    #[inline]
647    pub fn get(&self) -> T {
648        unsafe { *self.0.get() }
649    }
650
651    /// Returns a pointer to the current value.
652    #[inline]
653    pub fn as_ptr(&self) -> *const T {
654        self.0.get()
655    }
656}
657
658macro_rules! impl_element(
659    ($t:ty, $f:ident) => {
660        unsafe impl Element for $t {
661            fn is_compatible_format(format: &CStr) -> bool {
662                let slice = format.to_bytes();
663                if slice.len() > 1 && !is_matching_endian(slice[0]) {
664                    return false;
665                }
666                ElementType::from_format(format) == ElementType::$f { bytes: mem::size_of::<$t>() }
667            }
668        }
669    }
670);
671
672impl_element!(u8, UnsignedInteger);
673impl_element!(u16, UnsignedInteger);
674impl_element!(u32, UnsignedInteger);
675impl_element!(u64, UnsignedInteger);
676impl_element!(usize, UnsignedInteger);
677impl_element!(i8, SignedInteger);
678impl_element!(i16, SignedInteger);
679impl_element!(i32, SignedInteger);
680impl_element!(i64, SignedInteger);
681impl_element!(isize, SignedInteger);
682impl_element!(f32, Float);
683impl_element!(f64, Float);
684
685#[cfg(test)]
686mod tests {
687    use super::PyBuffer;
688    use crate::ffi;
689    use crate::types::any::PyAnyMethods;
690    use crate::Python;
691
692    #[test]
693    fn test_debug() {
694        Python::with_gil(|py| {
695            let bytes = py.eval(ffi::c_str!("b'abcde'"), None, None).unwrap();
696            let buffer: PyBuffer<u8> = PyBuffer::get(&bytes).unwrap();
697            let expected = format!(
698                concat!(
699                    "PyBuffer {{ buf: {:?}, obj: {:?}, ",
700                    "len: 5, itemsize: 1, readonly: 1, ",
701                    "ndim: 1, format: {:?}, shape: {:?}, ",
702                    "strides: {:?}, suboffsets: {:?}, internal: {:?} }}",
703                ),
704                buffer.0.buf,
705                buffer.0.obj,
706                buffer.0.format,
707                buffer.0.shape,
708                buffer.0.strides,
709                buffer.0.suboffsets,
710                buffer.0.internal
711            );
712            let debug_repr = format!("{:?}", buffer);
713            assert_eq!(debug_repr, expected);
714        });
715    }
716
717    #[test]
718    fn test_element_type_from_format() {
719        use super::ElementType;
720        use super::ElementType::*;
721        use std::mem::size_of;
722        use std::os::raw;
723
724        for (cstr, expected) in [
725            // @ prefix goes to native_element_type_from_type_char
726            (
727                ffi::c_str!("@b"),
728                SignedInteger {
729                    bytes: size_of::<raw::c_schar>(),
730                },
731            ),
732            (
733                ffi::c_str!("@c"),
734                UnsignedInteger {
735                    bytes: size_of::<raw::c_char>(),
736                },
737            ),
738            (
739                ffi::c_str!("@b"),
740                SignedInteger {
741                    bytes: size_of::<raw::c_schar>(),
742                },
743            ),
744            (
745                ffi::c_str!("@B"),
746                UnsignedInteger {
747                    bytes: size_of::<raw::c_uchar>(),
748                },
749            ),
750            (ffi::c_str!("@?"), Bool),
751            (
752                ffi::c_str!("@h"),
753                SignedInteger {
754                    bytes: size_of::<raw::c_short>(),
755                },
756            ),
757            (
758                ffi::c_str!("@H"),
759                UnsignedInteger {
760                    bytes: size_of::<raw::c_ushort>(),
761                },
762            ),
763            (
764                ffi::c_str!("@i"),
765                SignedInteger {
766                    bytes: size_of::<raw::c_int>(),
767                },
768            ),
769            (
770                ffi::c_str!("@I"),
771                UnsignedInteger {
772                    bytes: size_of::<raw::c_uint>(),
773                },
774            ),
775            (
776                ffi::c_str!("@l"),
777                SignedInteger {
778                    bytes: size_of::<raw::c_long>(),
779                },
780            ),
781            (
782                ffi::c_str!("@L"),
783                UnsignedInteger {
784                    bytes: size_of::<raw::c_ulong>(),
785                },
786            ),
787            (
788                ffi::c_str!("@q"),
789                SignedInteger {
790                    bytes: size_of::<raw::c_longlong>(),
791                },
792            ),
793            (
794                ffi::c_str!("@Q"),
795                UnsignedInteger {
796                    bytes: size_of::<raw::c_ulonglong>(),
797                },
798            ),
799            (
800                ffi::c_str!("@n"),
801                SignedInteger {
802                    bytes: size_of::<libc::ssize_t>(),
803                },
804            ),
805            (
806                ffi::c_str!("@N"),
807                UnsignedInteger {
808                    bytes: size_of::<libc::size_t>(),
809                },
810            ),
811            (ffi::c_str!("@e"), Float { bytes: 2 }),
812            (ffi::c_str!("@f"), Float { bytes: 4 }),
813            (ffi::c_str!("@d"), Float { bytes: 8 }),
814            (ffi::c_str!("@z"), Unknown),
815            // = prefix goes to standard_element_type_from_type_char
816            (ffi::c_str!("=b"), SignedInteger { bytes: 1 }),
817            (ffi::c_str!("=c"), UnsignedInteger { bytes: 1 }),
818            (ffi::c_str!("=B"), UnsignedInteger { bytes: 1 }),
819            (ffi::c_str!("=?"), Bool),
820            (ffi::c_str!("=h"), SignedInteger { bytes: 2 }),
821            (ffi::c_str!("=H"), UnsignedInteger { bytes: 2 }),
822            (ffi::c_str!("=l"), SignedInteger { bytes: 4 }),
823            (ffi::c_str!("=l"), SignedInteger { bytes: 4 }),
824            (ffi::c_str!("=I"), UnsignedInteger { bytes: 4 }),
825            (ffi::c_str!("=L"), UnsignedInteger { bytes: 4 }),
826            (ffi::c_str!("=q"), SignedInteger { bytes: 8 }),
827            (ffi::c_str!("=Q"), UnsignedInteger { bytes: 8 }),
828            (ffi::c_str!("=e"), Float { bytes: 2 }),
829            (ffi::c_str!("=f"), Float { bytes: 4 }),
830            (ffi::c_str!("=d"), Float { bytes: 8 }),
831            (ffi::c_str!("=z"), Unknown),
832            (ffi::c_str!("=0"), Unknown),
833            // unknown prefix -> Unknown
834            (ffi::c_str!(":b"), Unknown),
835        ] {
836            assert_eq!(
837                ElementType::from_format(cstr),
838                expected,
839                "element from format &Cstr: {:?}",
840                cstr,
841            );
842        }
843    }
844
845    #[test]
846    fn test_compatible_size() {
847        // for the cast in PyBuffer::shape()
848        assert_eq!(
849            std::mem::size_of::<ffi::Py_ssize_t>(),
850            std::mem::size_of::<usize>()
851        );
852    }
853
854    #[test]
855    fn test_bytes_buffer() {
856        Python::with_gil(|py| {
857            let bytes = py.eval(ffi::c_str!("b'abcde'"), None, None).unwrap();
858            let buffer = PyBuffer::get(&bytes).unwrap();
859            assert_eq!(buffer.dimensions(), 1);
860            assert_eq!(buffer.item_count(), 5);
861            assert_eq!(buffer.format().to_str().unwrap(), "B");
862            assert_eq!(buffer.shape(), [5]);
863            // single-dimensional buffer is always contiguous
864            assert!(buffer.is_c_contiguous());
865            assert!(buffer.is_fortran_contiguous());
866
867            let slice = buffer.as_slice(py).unwrap();
868            assert_eq!(slice.len(), 5);
869            assert_eq!(slice[0].get(), b'a');
870            assert_eq!(slice[2].get(), b'c');
871
872            assert_eq!(unsafe { *(buffer.get_ptr(&[1]) as *mut u8) }, b'b');
873
874            assert!(buffer.as_mut_slice(py).is_none());
875
876            assert!(buffer.copy_to_slice(py, &mut [0u8]).is_err());
877            let mut arr = [0; 5];
878            buffer.copy_to_slice(py, &mut arr).unwrap();
879            assert_eq!(arr, b"abcde" as &[u8]);
880
881            assert!(buffer.copy_from_slice(py, &[0u8; 5]).is_err());
882            assert_eq!(buffer.to_vec(py).unwrap(), b"abcde");
883        });
884    }
885
886    #[test]
887    fn test_array_buffer() {
888        Python::with_gil(|py| {
889            let array = py
890                .import("array")
891                .unwrap()
892                .call_method("array", ("f", (1.0, 1.5, 2.0, 2.5)), None)
893                .unwrap();
894            let buffer = PyBuffer::get(&array).unwrap();
895            assert_eq!(buffer.dimensions(), 1);
896            assert_eq!(buffer.item_count(), 4);
897            assert_eq!(buffer.format().to_str().unwrap(), "f");
898            assert_eq!(buffer.shape(), [4]);
899
900            // array creates a 1D contiguious buffer, so it's both C and F contiguous.  This would
901            // be more interesting if we can come up with a 2D buffer but I think it would need a
902            // third-party lib or a custom class.
903
904            // C-contiguous fns
905            let slice = buffer.as_slice(py).unwrap();
906            assert_eq!(slice.len(), 4);
907            assert_eq!(slice[0].get(), 1.0);
908            assert_eq!(slice[3].get(), 2.5);
909
910            let mut_slice = buffer.as_mut_slice(py).unwrap();
911            assert_eq!(mut_slice.len(), 4);
912            assert_eq!(mut_slice[0].get(), 1.0);
913            mut_slice[3].set(2.75);
914            assert_eq!(slice[3].get(), 2.75);
915
916            buffer
917                .copy_from_slice(py, &[10.0f32, 11.0, 12.0, 13.0])
918                .unwrap();
919            assert_eq!(slice[2].get(), 12.0);
920
921            assert_eq!(buffer.to_vec(py).unwrap(), [10.0, 11.0, 12.0, 13.0]);
922
923            // F-contiguous fns
924            let buffer = PyBuffer::get(&array).unwrap();
925            let slice = buffer.as_fortran_slice(py).unwrap();
926            assert_eq!(slice.len(), 4);
927            assert_eq!(slice[1].get(), 11.0);
928
929            let mut_slice = buffer.as_fortran_mut_slice(py).unwrap();
930            assert_eq!(mut_slice.len(), 4);
931            assert_eq!(mut_slice[2].get(), 12.0);
932            mut_slice[3].set(2.75);
933            assert_eq!(slice[3].get(), 2.75);
934
935            buffer
936                .copy_from_fortran_slice(py, &[10.0f32, 11.0, 12.0, 13.0])
937                .unwrap();
938            assert_eq!(slice[2].get(), 12.0);
939
940            assert_eq!(buffer.to_fortran_vec(py).unwrap(), [10.0, 11.0, 12.0, 13.0]);
941        });
942    }
943}
⚠️ Internal Docs ⚠️ Not Public API 👉 Official Docs Here