pyo3/conversions/std/
array.rs

1// TODO https://github.com/PyO3/pyo3/issues/5487
2#![allow(clippy::undocumented_unsafe_blocks)]
3
4use crate::conversion::{FromPyObjectOwned, FromPyObjectSequence, IntoPyObject};
5#[cfg(feature = "experimental-inspect")]
6use crate::inspect::{type_hint_subscript, PyStaticExpr};
7use crate::types::any::PyAnyMethods;
8use crate::types::PySequence;
9use crate::{err::CastError, ffi, FromPyObject, PyAny, PyResult, PyTypeInfo, Python};
10use crate::{exceptions, Borrowed, Bound, PyErr};
11
12impl<'py, T, const N: usize> IntoPyObject<'py> for [T; N]
13where
14    T: IntoPyObject<'py>,
15{
16    type Target = PyAny;
17    type Output = Bound<'py, Self::Target>;
18    type Error = PyErr;
19
20    #[cfg(feature = "experimental-inspect")]
21    const OUTPUT_TYPE: PyStaticExpr = T::SEQUENCE_OUTPUT_TYPE;
22
23    /// Turns [`[u8; N]`](core::array) into [`PyBytes`], all other `T`s will be turned into a [`PyList`]
24    ///
25    /// [`PyBytes`]: crate::types::PyBytes
26    /// [`PyList`]: crate::types::PyList
27    #[inline]
28    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
29        T::owned_sequence_into_pyobject(self, py, crate::conversion::private::Token)
30    }
31}
32
33impl<'a, 'py, T, const N: usize> IntoPyObject<'py> for &'a [T; N]
34where
35    &'a T: IntoPyObject<'py>,
36{
37    type Target = PyAny;
38    type Output = Bound<'py, Self::Target>;
39    type Error = PyErr;
40
41    #[cfg(feature = "experimental-inspect")]
42    const OUTPUT_TYPE: PyStaticExpr = <&[T]>::OUTPUT_TYPE;
43
44    #[inline]
45    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
46        self.as_slice().into_pyobject(py)
47    }
48}
49
50impl<'py, T, const N: usize> FromPyObject<'_, 'py> for [T; N]
51where
52    T: FromPyObjectOwned<'py>,
53{
54    type Error = PyErr;
55
56    #[cfg(feature = "experimental-inspect")]
57    const INPUT_TYPE: PyStaticExpr = type_hint_subscript!(PySequence::TYPE_HINT, T::INPUT_TYPE);
58
59    fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {
60        if let Some(extractor) = T::sequence_extractor(obj, crate::conversion::private::Token) {
61            return extractor.to_array();
62        }
63
64        create_array_from_obj(obj)
65    }
66}
67
68fn create_array_from_obj<'py, T, const N: usize>(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<[T; N]>
69where
70    T: FromPyObjectOwned<'py>,
71{
72    // Types that pass `PySequence_Check` usually implement enough of the sequence protocol
73    // to support this function and if not, we will only fail extraction safely.
74    if unsafe { ffi::PySequence_Check(obj.as_ptr()) } == 0 {
75        return Err(CastError::new(obj, PySequence::type_object(obj.py()).into_any()).into());
76    }
77
78    let seq_len = obj.len()?;
79    if seq_len != N {
80        return Err(invalid_sequence_length(N, seq_len));
81    }
82    array_try_from_fn(|idx| {
83        obj.get_item(idx)
84            .and_then(|any| any.extract().map_err(Into::into))
85    })
86}
87
88// TODO use core::array::try_from_fn, if that stabilises:
89// (https://github.com/rust-lang/rust/issues/89379)
90fn array_try_from_fn<E, F, T, const N: usize>(mut cb: F) -> Result<[T; N], E>
91where
92    F: FnMut(usize) -> Result<T, E>,
93{
94    // Helper to safely create arrays since the standard library doesn't
95    // provide one yet. Shouldn't be necessary in the future.
96    struct ArrayGuard<T, const N: usize> {
97        dst: *mut T,
98        initialized: usize,
99    }
100
101    impl<T, const N: usize> Drop for ArrayGuard<T, N> {
102        fn drop(&mut self) {
103            debug_assert!(self.initialized <= N);
104            let initialized_part = core::ptr::slice_from_raw_parts_mut(self.dst, self.initialized);
105            unsafe {
106                core::ptr::drop_in_place(initialized_part);
107            }
108        }
109    }
110
111    // [MaybeUninit<T>; N] would be "nicer" but is actually difficult to create - there are nightly
112    // APIs which would make this easier.
113    let mut array: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit();
114    let mut guard: ArrayGuard<T, N> = ArrayGuard {
115        dst: array.as_mut_ptr() as _,
116        initialized: 0,
117    };
118    unsafe {
119        let mut value_ptr = array.as_mut_ptr() as *mut T;
120        for i in 0..N {
121            core::ptr::write(value_ptr, cb(i)?);
122            value_ptr = value_ptr.offset(1);
123            guard.initialized += 1;
124        }
125        core::mem::forget(guard);
126        Ok(array.assume_init())
127    }
128}
129
130pub(crate) fn invalid_sequence_length(expected: usize, actual: usize) -> PyErr {
131    exceptions::PyValueError::new_err(format!(
132        "expected a sequence of length {expected} (got {actual})"
133    ))
134}
135
136#[cfg(test)]
137mod tests {
138    #[cfg(panic = "unwind")]
139    use core::sync::atomic::{AtomicUsize, Ordering};
140    #[cfg(panic = "unwind")]
141    use std::panic;
142
143    use crate::{
144        conversion::IntoPyObject,
145        types::{any::PyAnyMethods, PyBytes, PyBytesMethods},
146    };
147    use crate::{types::PyList, PyResult, Python};
148
149    #[test]
150    #[cfg(panic = "unwind")]
151    fn array_try_from_fn() {
152        static DROP_COUNTER: AtomicUsize = AtomicUsize::new(0);
153        struct CountDrop;
154        impl Drop for CountDrop {
155            fn drop(&mut self) {
156                DROP_COUNTER.fetch_add(1, Ordering::SeqCst);
157            }
158        }
159        let _ = catch_unwind_silent(move || {
160            let _: Result<[CountDrop; 4], ()> = super::array_try_from_fn(|idx| {
161                #[expect(clippy::manual_assert, reason = "testing panic during array creation")]
162                if idx == 2 {
163                    panic!("peek a boo");
164                }
165                Ok(CountDrop)
166            });
167        });
168        assert_eq!(DROP_COUNTER.load(Ordering::SeqCst), 2);
169    }
170
171    #[test]
172    fn test_extract_bytes_to_array() {
173        Python::attach(|py| {
174            let v: [u8; 33] = py
175                .eval(c"b'abcabcabcabcabcabcabcabcabcabcabc'", None, None)
176                .unwrap()
177                .extract()
178                .unwrap();
179            assert_eq!(&v, b"abcabcabcabcabcabcabcabcabcabcabc");
180        })
181    }
182
183    #[test]
184    fn test_extract_bytes_wrong_length() {
185        Python::attach(|py| {
186            let v: PyResult<[u8; 3]> = py.eval(c"b'abcdefg'", None, None).unwrap().extract();
187            assert_eq!(
188                v.unwrap_err().to_string(),
189                "ValueError: expected a sequence of length 3 (got 7)"
190            );
191        })
192    }
193
194    #[test]
195    fn test_extract_bytearray_to_array() {
196        Python::attach(|py| {
197            let v: [u8; 33] = py
198                .eval(
199                    c"bytearray(b'abcabcabcabcabcabcabcabcabcabcabc')",
200                    None,
201                    None,
202                )
203                .unwrap()
204                .extract()
205                .unwrap();
206            assert_eq!(&v, b"abcabcabcabcabcabcabcabcabcabcabc");
207        })
208    }
209
210    #[test]
211    fn test_extract_small_bytearray_to_array() {
212        Python::attach(|py| {
213            let v: [u8; 3] = py
214                .eval(c"bytearray(b'abc')", None, None)
215                .unwrap()
216                .extract()
217                .unwrap();
218            assert_eq!(&v, b"abc");
219        });
220    }
221    #[test]
222    fn test_into_pyobject_array_conversion() {
223        Python::attach(|py| {
224            let array: [f32; 4] = [0.0, -16.0, 16.0, 42.0];
225            let pyobject = array.into_pyobject(py).unwrap();
226            let pylist = pyobject.cast::<PyList>().unwrap();
227            assert_eq!(pylist.get_item(0).unwrap().extract::<f32>().unwrap(), 0.0);
228            assert_eq!(pylist.get_item(1).unwrap().extract::<f32>().unwrap(), -16.0);
229            assert_eq!(pylist.get_item(2).unwrap().extract::<f32>().unwrap(), 16.0);
230            assert_eq!(pylist.get_item(3).unwrap().extract::<f32>().unwrap(), 42.0);
231        });
232    }
233
234    #[test]
235    fn test_extract_invalid_sequence_length() {
236        Python::attach(|py| {
237            let v: PyResult<[u8; 3]> = py
238                .eval(c"bytearray(b'abcdefg')", None, None)
239                .unwrap()
240                .extract();
241            assert_eq!(
242                v.unwrap_err().to_string(),
243                "ValueError: expected a sequence of length 3 (got 7)"
244            );
245        })
246    }
247
248    #[test]
249    fn test_intopyobject_array_conversion() {
250        Python::attach(|py| {
251            let array: [f32; 4] = [0.0, -16.0, 16.0, 42.0];
252            let pylist = array
253                .into_pyobject(py)
254                .unwrap()
255                .cast_into::<PyList>()
256                .unwrap();
257
258            assert_eq!(pylist.get_item(0).unwrap().extract::<f32>().unwrap(), 0.0);
259            assert_eq!(pylist.get_item(1).unwrap().extract::<f32>().unwrap(), -16.0);
260            assert_eq!(pylist.get_item(2).unwrap().extract::<f32>().unwrap(), 16.0);
261            assert_eq!(pylist.get_item(3).unwrap().extract::<f32>().unwrap(), 42.0);
262        });
263    }
264
265    #[test]
266    fn test_array_intopyobject_impl() {
267        Python::attach(|py| {
268            let bytes: [u8; 6] = *b"foobar";
269            let obj = bytes.into_pyobject(py).unwrap();
270            assert!(obj.is_instance_of::<PyBytes>());
271            let obj = obj.cast_into::<PyBytes>().unwrap();
272            assert_eq!(obj.as_bytes(), &bytes);
273
274            let nums: [u16; 4] = [0, 1, 2, 3];
275            let obj = nums.into_pyobject(py).unwrap();
276            assert!(obj.is_instance_of::<PyList>());
277        });
278    }
279
280    #[test]
281    fn test_extract_non_iterable_to_array() {
282        Python::attach(|py| {
283            let v = py.eval(c"42", None, None).unwrap();
284            v.extract::<i32>().unwrap();
285            v.extract::<[i32; 1]>().unwrap_err();
286        });
287    }
288
289    #[cfg(feature = "macros")]
290    #[test]
291    fn test_pyclass_intopy_array_conversion() {
292        #[crate::pyclass(crate = "crate")]
293        struct Foo;
294
295        Python::attach(|py| {
296            let array: [Foo; 8] = [Foo, Foo, Foo, Foo, Foo, Foo, Foo, Foo];
297            let list = array
298                .into_pyobject(py)
299                .unwrap()
300                .cast_into::<PyList>()
301                .unwrap();
302            let _bound = list.get_item(4).unwrap().cast::<Foo>().unwrap();
303        });
304    }
305
306    // https://stackoverflow.com/a/59211505
307    #[cfg(panic = "unwind")]
308    fn catch_unwind_silent<F, R>(f: F) -> std::thread::Result<R>
309    where
310        F: FnOnce() -> R + panic::UnwindSafe,
311    {
312        let prev_hook = panic::take_hook();
313        panic::set_hook(Box::new(|_| {}));
314        let result = panic::catch_unwind(f);
315        panic::set_hook(prev_hook);
316        result
317    }
318}
pyo3/conversions/std/array.rs

pyo3/conversions/std/
array.rs
