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