From 9bc5089e4b69da220a51ec9f177f3e873c365802 Mon Sep 17 00:00:00 2001 From: Caio Date: Wed, 31 Mar 2021 12:04:15 -0300 Subject: [PATCH] Add support for arbitrary arrays --- CHANGELOG.md | 1 + build.rs | 17 +++ src/conversions/array.rs | 273 +++++++++++++++++++++++++++++++++++++++ src/conversions/mod.rs | 1 + src/lib.rs | 1 + src/types/list.rs | 20 --- src/types/sequence.rs | 82 ------------ src/utils.rs | 8 ++ 8 files changed, 301 insertions(+), 102 deletions(-) create mode 100644 src/conversions/array.rs create mode 100644 src/utils.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e75f883..6c6fb353 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. - Support PyPy 3.7. [#1538](https://github.com/PyO3/pyo3/pull/1538) ### Added +- Add conversion for `[T; N]` for all `N` on Rust 1.51 and up. [#1128](https://github.com/PyO3/pyo3/pull/1128) - Add conversions between `OsStr`/`OsString`/`Path`/`PathBuf` and Python strings. [#1379](https://github.com/PyO3/pyo3/pull/1379) - Add `#[pyo3(from_py_with = "...")]` attribute for function arguments and struct fields to override the default from-Python conversion. [#1411](https://github.com/PyO3/pyo3/pull/1411) - Add FFI definition `PyCFunction_CheckExact` for Python 3.9 and later. [#1425](https://github.com/PyO3/pyo3/pull/1425) diff --git a/build.rs b/build.rs index 55e4b08b..c67e92bc 100644 --- a/build.rs +++ b/build.rs @@ -746,6 +746,17 @@ fn ensure_python_version(interpreter_config: &InterpreterConfig) -> Result<()> { Ok(()) } +fn rustc_minor_version() -> Option { + let rustc = env::var_os("RUSTC")?; + let output = Command::new(rustc).arg("--version").output().ok()?; + let version = core::str::from_utf8(&output.stdout).ok()?; + let mut pieces = version.split('.'); + if pieces.next() != Some("rustc 1") { + return None; + } + pieces.next()?.parse().ok() +} + fn emit_cargo_configuration(interpreter_config: &InterpreterConfig) -> Result<()> { let target_os = cargo_env_var("CARGO_CFG_TARGET_OS").unwrap(); let is_extension_module = cargo_env_var("CARGO_FEATURE_EXTENSION_MODULE").is_some(); @@ -850,6 +861,12 @@ fn emit_cargo_configuration(interpreter_config: &InterpreterConfig) -> Result<() println!("cargo:rustc-cfg=py_sys_config=\"{}\"", flag) } + // Enable use of const generics on Rust 1.51 and greater + + if rustc_minor_version().unwrap_or(0) >= 51 { + println!("cargo:rustc-cfg=min_const_generics"); + } + Ok(()) } diff --git a/src/conversions/array.rs b/src/conversions/array.rs new file mode 100644 index 00000000..8f6f321d --- /dev/null +++ b/src/conversions/array.rs @@ -0,0 +1,273 @@ +use crate::{FromPyObject, IntoPy, PyAny, PyObject, PyResult, PyTryFrom, Python, ToPyObject}; + +#[cfg(not(min_const_generics))] +macro_rules! array_impls { + ($($N:expr),+) => { + $( + impl<'a, T> FromPyObject<'a> for [T; $N] + where + T: Copy + Default + FromPyObject<'a>, + { + #[cfg(not(feature = "nightly"))] + fn extract(obj: &'a PyAny) -> PyResult { + let mut array = [T::default(); $N]; + _extract_sequence_into_slice(obj, &mut array)?; + Ok(array) + } + + #[cfg(feature = "nightly")] + default fn extract(obj: &'a PyAny) -> PyResult { + let mut array = [T::default(); $N]; + _extract_sequence_into_slice(obj, &mut array)?; + Ok(array) + } + } + + #[cfg(feature = "nightly")] + impl<'source, T> FromPyObject<'source> for [T; $N] + where + for<'a> T: Default + FromPyObject<'a> + crate::buffer::Element, + { + fn extract(obj: &'source PyAny) -> PyResult { + let mut array = [T::default(); $N]; + // first try buffer protocol + if unsafe { crate::ffi::PyObject_CheckBuffer(obj.as_ptr()) } == 1 { + if let Ok(buf) = crate::buffer::PyBuffer::get(obj) { + if buf.dimensions() == 1 && buf.copy_to_slice(obj.py(), &mut array).is_ok() { + buf.release(obj.py()); + return Ok(array); + } + buf.release(obj.py()); + } + } + // fall back to sequence protocol + _extract_sequence_into_slice(obj, &mut array)?; + Ok(array) + } + } + )+ + } +} + +#[cfg(not(min_const_generics))] +array_impls!( + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, + 26, 27, 28, 29, 30, 31, 32 +); + +#[cfg(min_const_generics)] +impl<'a, T, const N: usize> FromPyObject<'a> for [T; N] +where + T: FromPyObject<'a>, +{ + #[cfg(not(feature = "nightly"))] + fn extract(obj: &'a PyAny) -> PyResult { + create_array_from_obj(obj) + } + + #[cfg(feature = "nightly")] + default fn extract(obj: &'a PyAny) -> PyResult { + create_array_from_obj(obj) + } +} + +#[cfg(not(min_const_generics))] +macro_rules! array_impls { + ($($N:expr),+) => { + $( + impl IntoPy for [T; $N] + where + T: ToPyObject + { + fn into_py(self, py: Python) -> PyObject { + self.as_ref().to_object(py) + } + } + )+ + } +} + +#[cfg(not(min_const_generics))] +array_impls!( + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, + 26, 27, 28, 29, 30, 31, 32 +); + +#[cfg(min_const_generics)] +impl IntoPy for [T; N] +where + T: ToPyObject, +{ + fn into_py(self, py: Python) -> PyObject { + self.as_ref().to_object(py) + } +} + +#[cfg(all(min_const_generics, feature = "nightly"))] +impl<'source, T, const N: usize> FromPyObject<'source> for [T; N] +where + for<'a> T: FromPyObject<'a> + crate::buffer::Element, +{ + fn extract(obj: &'source PyAny) -> PyResult { + let mut array: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit(); + // first try buffer protocol + if unsafe { crate::ffi::PyObject_CheckBuffer(obj.as_ptr()) } == 1 { + if let Ok(buf) = crate::buffer::PyBuffer::get(obj) { + if buf.dimensions() == 1 && buf.copy_to_slice(obj.py(), &mut array).is_ok() { + buf.release(obj.py()); + // SAFETY: The array should be fully filled by `copy_to_slice` + return Ok(unsafe { array.assume_init() }); + } + buf.release(obj.py()); + } + } + // fall back to sequence protocol + _extract_sequence_into_slice(obj, &mut array)?; + // SAFETY: The array should be fully filled by `_extract_sequence_into_slice` + Ok(unsafe { array.assume_init() }) + } +} + +#[cfg(min_const_generics)] +fn create_array_from_obj<'s, T, const N: usize>(obj: &'s PyAny) -> PyResult<[T; N]> +where + T: FromPyObject<'s>, +{ + let seq = ::try_from(obj)?; + let expected_len = seq.len()? as usize; + let mut counter = 0; + try_create_array(&mut counter, |idx| { + seq.get_item(idx as isize) + .map_err(|_| crate::utils::invalid_sequence_length(expected_len, idx + 1))? + .extract::() + }) +} + +fn _extract_sequence_into_slice<'s, T>(obj: &'s PyAny, slice: &mut [T]) -> PyResult<()> +where + T: FromPyObject<'s>, +{ + let seq = ::try_from(obj)?; + let expected_len = seq.len()? as usize; + if expected_len != slice.len() { + return Err(crate::utils::invalid_sequence_length( + expected_len, + slice.len(), + )); + } + for (value, item) in slice.iter_mut().zip(seq.iter()?) { + *value = item?.extract::()?; + } + Ok(()) +} + +#[cfg(min_const_generics)] +fn try_create_array(counter: &mut usize, mut cb: F) -> Result<[T; N], E> +where + F: FnMut(usize) -> Result, +{ + // Helper to safely create arrays since the standard library doesn't + // provide one yet. Shouldn't be necessary in the future. + struct ArrayGuard<'a, T, const N: usize> { + dst: *mut T, + initialized: &'a mut usize, + } + + impl Drop for ArrayGuard<'_, T, N> { + fn drop(&mut self) { + debug_assert!(*self.initialized <= N); + let initialized_part = core::ptr::slice_from_raw_parts_mut(self.dst, *self.initialized); + unsafe { + core::ptr::drop_in_place(initialized_part); + } + } + } + + let mut array: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit(); + let guard: ArrayGuard = ArrayGuard { + dst: array.as_mut_ptr() as _, + initialized: counter, + }; + unsafe { + for (idx, value_ptr) in (&mut *array.as_mut_ptr()).iter_mut().enumerate() { + core::ptr::write(value_ptr, cb(idx)?); + *guard.initialized += 1; + } + core::mem::forget(guard); + Ok(array.assume_init()) + } +} + +#[cfg(test)] +mod test { + use crate::Python; + #[cfg(min_const_generics)] + use std::{ + panic, + sync::{Arc, Mutex}, + thread::sleep, + time, + }; + + #[cfg(min_const_generics)] + #[test] + fn try_create_array() { + #[allow(clippy::mutex_atomic)] + let counter = Arc::new(Mutex::new(0)); + let counter_unwind = Arc::clone(&counter); + let _ = catch_unwind_silent(move || { + let mut locked = counter_unwind.lock().unwrap(); + let _: Result<[i32; 4], _> = super::try_create_array(&mut *locked, |idx| { + if idx == 2 { + panic!("peek a boo"); + } + Ok::<_, ()>(1) + }); + }); + sleep(time::Duration::from_secs(2)); + assert_eq!(*counter.lock().unwrap_err().into_inner(), 2); + } + + #[cfg(not(min_const_generics))] + #[test] + fn test_extract_bytearray_to_array() { + let gil = Python::acquire_gil(); + let py = gil.python(); + let v: [u8; 3] = py + .eval("bytearray(b'abc')", None, None) + .unwrap() + .extract() + .unwrap(); + assert!(&v == b"abc"); + } + + #[cfg(min_const_generics)] + #[test] + fn test_extract_bytearray_to_array() { + let gil = Python::acquire_gil(); + let py = gil.python(); + let v: [u8; 33] = py + .eval( + "bytearray(b'abcabcabcabcabcabcabcabcabcabcabc')", + None, + None, + ) + .unwrap() + .extract() + .unwrap(); + assert!(&v == b"abcabcabcabcabcabcabcabcabcabcabc"); + } + + // https://stackoverflow.com/a/59211505 + #[cfg(min_const_generics)] + fn catch_unwind_silent(f: F) -> std::thread::Result + where + F: FnOnce() -> R + panic::UnwindSafe, + { + let prev_hook = panic::take_hook(); + panic::set_hook(Box::new(|_| {})); + let result = panic::catch_unwind(f); + panic::set_hook(prev_hook); + result + } +} diff --git a/src/conversions/mod.rs b/src/conversions/mod.rs index a7ccff99..9be3ba9f 100644 --- a/src/conversions/mod.rs +++ b/src/conversions/mod.rs @@ -1,5 +1,6 @@ //! This module contains conversions between non-String Rust object and their string representation //! in Python +mod array; mod osstr; mod path; diff --git a/src/lib.rs b/src/lib.rs index 69735f74..f3375910 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -186,6 +186,7 @@ pub mod pyclass_slots; mod python; pub mod type_object; pub mod types; +mod utils; #[cfg(feature = "serde")] pub mod serde; diff --git a/src/types/list.rs b/src/types/list.rs index 4a0586e4..019693b1 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -178,26 +178,6 @@ where } } -macro_rules! array_impls { - ($($N:expr),+) => { - $( - impl IntoPy for [T; $N] - where - T: ToPyObject - { - fn into_py(self, py: Python) -> PyObject { - self.as_ref().to_object(py) - } - } - )+ - } -} - -array_impls!( - 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, - 26, 27, 28, 29, 30, 31, 32 -); - impl ToPyObject for Vec where T: ToPyObject, diff --git a/src/types/sequence.rs b/src/types/sequence.rs index 423021d0..6432c4a0 100644 --- a/src/types/sequence.rs +++ b/src/types/sequence.rs @@ -1,7 +1,6 @@ // Copyright (c) 2017-present PyO3 Project and Contributors use crate::err::{self, PyDowncastError, PyErr, PyResult}; -use crate::exceptions; use crate::ffi::{self, Py_ssize_t}; use crate::instance::PyNativeType; use crate::types::{PyAny, PyList, PyTuple}; @@ -257,59 +256,6 @@ impl PySequence { } } -macro_rules! array_impls { - ($($N:expr),+) => { - $( - impl<'a, T> FromPyObject<'a> for [T; $N] - where - T: Copy + Default + FromPyObject<'a>, - { - #[cfg(not(feature = "nightly"))] - fn extract(obj: &'a PyAny) -> PyResult { - let mut array = [T::default(); $N]; - extract_sequence_into_slice(obj, &mut array)?; - Ok(array) - } - - #[cfg(feature = "nightly")] - default fn extract(obj: &'a PyAny) -> PyResult { - let mut array = [T::default(); $N]; - extract_sequence_into_slice(obj, &mut array)?; - Ok(array) - } - } - - #[cfg(feature = "nightly")] - impl<'source, T> FromPyObject<'source> for [T; $N] - where - for<'a> T: Default + FromPyObject<'a> + crate::buffer::Element, - { - fn extract(obj: &'source PyAny) -> PyResult { - let mut array = [T::default(); $N]; - // first try buffer protocol - if unsafe { ffi::PyObject_CheckBuffer(obj.as_ptr()) } == 1 { - if let Ok(buf) = crate::buffer::PyBuffer::get(obj) { - if buf.dimensions() == 1 && buf.copy_to_slice(obj.py(), &mut array).is_ok() { - buf.release(obj.py()); - return Ok(array); - } - buf.release(obj.py()); - } - } - // fall back to sequence protocol - extract_sequence_into_slice(obj, &mut array)?; - Ok(array) - } - } - )+ - } -} - -array_impls!( - 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, - 26, 27, 28, 29, 30, 31, 32 -); - impl<'a, T> FromPyObject<'a> for Vec where T: FromPyObject<'a>, @@ -357,22 +303,6 @@ where Ok(v) } -fn extract_sequence_into_slice<'s, T>(obj: &'s PyAny, slice: &mut [T]) -> PyResult<()> -where - T: FromPyObject<'s>, -{ - let seq = ::try_from(obj)?; - if seq.len()? as usize != slice.len() { - return Err(exceptions::PyBufferError::new_err( - "Slice length does not match buffer length.", - )); - } - for (value, item) in slice.iter_mut().zip(seq.iter()?) { - *value = item?.extract::()?; - } - Ok(()) -} - impl<'v> PyTryFrom<'v> for PySequence { fn try_from>(value: V) -> Result<&'v PySequence, PyDowncastError<'v>> { let value = value.into(); @@ -706,18 +636,6 @@ mod test { assert!(v == [1, 2, 3, 4]); } - #[test] - fn test_extract_bytearray_to_array() { - let gil = Python::acquire_gil(); - let py = gil.python(); - let v: [u8; 3] = py - .eval("bytearray(b'abc')", None, None) - .unwrap() - .extract() - .unwrap(); - assert!(&v == b"abc"); - } - #[test] fn test_extract_bytearray_to_vec() { let gil = Python::acquire_gil(); diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 00000000..82c629fc --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,8 @@ +use crate::{exceptions, PyErr}; + +pub fn invalid_sequence_length(expected: usize, actual: usize) -> PyErr { + exceptions::PyValueError::new_err(format!( + "expected a sequence of length {} (got {})", + expected, actual + )) +}