1 | #![cfg (any(not(Py_LIMITED_API), Py_3_11))] |
2 | // Copyright (c) 2017 Daniel Grunwald |
3 | // |
4 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this |
5 | // software and associated documentation files (the "Software"), to deal in the Software |
6 | // without restriction, including without limitation the rights to use, copy, modify, merge, |
7 | // publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons |
8 | // to whom the Software is furnished to do so, subject to the following conditions: |
9 | // |
10 | // The above copyright notice and this permission notice shall be included in all copies or |
11 | // substantial portions of the Software. |
12 | // |
13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, |
14 | // INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR |
15 | // PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE |
16 | // FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR |
17 | // OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER |
18 | // DEALINGS IN THE SOFTWARE. |
19 | |
20 | //! `PyBuffer` implementation |
21 | use crate::{err, exceptions::PyBufferError, ffi, FromPyObject, PyAny, PyResult, Python}; |
22 | use std::marker::PhantomData; |
23 | use std::os::raw; |
24 | use std::pin::Pin; |
25 | use std::{cell, mem, ptr, slice}; |
26 | use 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)] |
31 | pub 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. |
35 | unsafe impl<T> Send for PyBuffer<T> {} |
36 | unsafe impl<T> Sync for PyBuffer<T> {} |
37 | |
38 | impl<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)] |
58 | pub 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 | |
80 | impl 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 | |
94 | fn 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 | |
144 | fn 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" )] |
164 | fn is_matching_endian(c: u8) -> bool { |
165 | c == b'@' || c == b'=' || c == b'>' |
166 | } |
167 | |
168 | #[cfg (target_endian = "big" )] |
169 | fn 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. |
178 | pub 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 | |
184 | impl<'source, T: Element> FromPyObject<'source> for PyBuffer<T> { |
185 | fn extract(obj: &PyAny) -> PyResult<PyBuffer<T>> { |
186 | Self::get(obj) |
187 | } |
188 | } |
189 | |
190 | impl<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 | |
632 | impl<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)] |
644 | pub struct ReadOnlyCell<T: Element>(cell::UnsafeCell<T>); |
645 | |
646 | impl<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 | |
660 | macro_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 | |
674 | impl_element!(u8, UnsignedInteger); |
675 | impl_element!(u16, UnsignedInteger); |
676 | impl_element!(u32, UnsignedInteger); |
677 | impl_element!(u64, UnsignedInteger); |
678 | impl_element!(usize, UnsignedInteger); |
679 | impl_element!(i8, SignedInteger); |
680 | impl_element!(i16, SignedInteger); |
681 | impl_element!(i32, SignedInteger); |
682 | impl_element!(i64, SignedInteger); |
683 | impl_element!(isize, SignedInteger); |
684 | impl_element!(f32, Float); |
685 | impl_element!(f64, Float); |
686 | |
687 | #[cfg (test)] |
688 | mod 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 | |