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

pyo3/
buffer.rs
