Skip to content

Commit

Permalink
Implement *Py* traits for arrays (up to 32)
Browse files Browse the repository at this point in the history
  • Loading branch information
c410-f3r committed Mar 2, 2020
1 parent 3ded0b2 commit 6b78f1b
Show file tree
Hide file tree
Showing 3 changed files with 103 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
* `FromPyObject` is now automatically implemented for `T: Clone` pyclasses. [#730](https://github.com/PyO3/pyo3/pull/730)
* `#[pyo3(get)]` and `#[pyo3(set)]` will now use the Rust doc-comment from the field for the Python property. [#755](https://github.com/PyO3/pyo3/pull/755)
* `#[setter]` functions may now take an argument of `Pyo3::Python`. [#760](https://github.com/PyO3/pyo3/pull/760)
* Implement `FromPyObject` and `IntoPy<PyObject>` traits for arrays (up to 32). [#778](https://github.com/PyO3/pyo3/pull/778)

### Fixed

Expand Down
32 changes: 31 additions & 1 deletion src/types/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,26 @@ where
}
}

macro_rules! array_impls {
($($N:expr),+) => {
$(
impl<T> IntoPy<PyObject> 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<T> ToPyObject for Vec<T>
where
T: ToPyObject,
Expand Down Expand Up @@ -209,7 +229,7 @@ mod test {
use crate::objectprotocol::ObjectProtocol;
use crate::types::PyList;
use crate::Python;
use crate::{PyTryFrom, ToPyObject};
use crate::{PyTryFrom, PyObject, IntoPy, ToPyObject};

#[test]
fn test_new() {
Expand Down Expand Up @@ -423,4 +443,14 @@ mod test {
assert_eq!(3, list.get_item(2).extract::<i32>().unwrap());
assert_eq!(2, list.get_item(3).extract::<i32>().unwrap());
}

#[test]
fn test_array_into_py() {
let gil = Python::acquire_gil();
let py = gil.python();
let array: PyObject = [1, 2].into_py(py);
let list = <PyList as PyTryFrom>::try_from(array.as_ref(py)).unwrap();
assert_eq!(1, list.get_item(0).extract::<i32>().unwrap());
assert_eq!(2, list.get_item(1).extract::<i32>().unwrap());
}
}
71 changes: 71 additions & 0 deletions src/types/sequence.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

use crate::buffer;
use crate::err::{self, PyDowncastError, PyErr, PyResult};
use crate::exceptions;
use crate::ffi::{self, Py_ssize_t};
use crate::instance::PyNativeType;
use crate::internal_tricks::Unsendable;
Expand Down Expand Up @@ -241,6 +242,48 @@ impl PySequence {
}
}

macro_rules! array_impls {
($($N:expr),+) => {
$(
impl<'a, T> FromPyObject<'a> for [T; $N]
where
T: Copy + Default + FromPyObject<'a>,
{
default fn extract(obj: &'a PyAny) -> PyResult<Self> {
let mut array = [T::default(); $N];
extract_sequence_into_slice(obj, &mut array)?;
Ok(array)
}
}

impl<'source, T> FromPyObject<'source> for [T; $N]
where
for<'a> T: Copy + Default + FromPyObject<'a> + buffer::Element,
{
fn extract(obj: &'source PyAny) -> PyResult<Self> {
let mut array = [T::default(); $N];
// first try buffer protocol
if let Ok(buf) = buffer::PyBuffer::get(obj.py(), 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<T>
where
T: FromPyObject<'a>,
Expand Down Expand Up @@ -282,6 +325,22 @@ where
Ok(v)
}

fn extract_sequence_into_slice<'s, T>(obj: &'s PyAny, slice: &mut [T]) -> PyResult<()>
where
T: FromPyObject<'s>,
{
let seq = <PySequence as PyTryFrom>::try_from(obj)?;
if seq.len()? as usize != slice.len() {
return Err(exceptions::BufferError::py_err(
"Slice length does not match buffer length.",
));
}
for (value, item) in slice.iter_mut().zip(seq.iter()?) {
*value = item?.extract::<T>()?;
}
Ok(())
}

impl<'v> PyTryFrom<'v> for PySequence {
fn try_from<V: Into<&'v PyAny>>(value: V) -> Result<&'v PySequence, PyDowncastError> {
let value = value.into();
Expand Down Expand Up @@ -640,6 +699,18 @@ 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();
Expand Down

0 comments on commit 6b78f1b

Please sign in to comment.