diff --git a/newsfragments/4667.added.md b/newsfragments/4667.added.md new file mode 100644 index 00000000000..fc2a914607e --- /dev/null +++ b/newsfragments/4667.added.md @@ -0,0 +1 @@ +Add `PyList_Extend` & `PyList_Clear` to pyo3-ffi diff --git a/pyo3-ffi/src/compat/py_3_13.rs b/pyo3-ffi/src/compat/py_3_13.rs index 9f44ced6f3f..59289cb76ae 100644 --- a/pyo3-ffi/src/compat/py_3_13.rs +++ b/pyo3-ffi/src/compat/py_3_13.rs @@ -83,3 +83,24 @@ compat_function!( 1 } ); + +compat_function!( + originally_defined_for(Py_3_13); + + #[inline] + pub unsafe fn PyList_Extend( + list: *mut crate::PyObject, + iterable: *mut crate::PyObject, + ) -> std::os::raw::c_int { + crate::PyList_SetSlice(list, crate::PY_SSIZE_T_MAX, crate::PY_SSIZE_T_MAX, iterable) + } +); + +compat_function!( + originally_defined_for(Py_3_13); + + #[inline] + pub unsafe fn PyList_Clear(list: *mut crate::PyObject) -> std::os::raw::c_int { + crate::PyList_SetSlice(list, 0, crate::PY_SSIZE_T_MAX, std::ptr::null_mut()) + } +); diff --git a/pyo3-ffi/src/listobject.rs b/pyo3-ffi/src/listobject.rs index 9d8b7ed6a58..881a8a8707b 100644 --- a/pyo3-ffi/src/listobject.rs +++ b/pyo3-ffi/src/listobject.rs @@ -50,6 +50,10 @@ extern "C" { arg3: Py_ssize_t, arg4: *mut PyObject, ) -> c_int; + #[cfg(Py_3_13)] + pub fn PyList_Extend(list: *mut PyObject, iterable: *mut PyObject) -> c_int; + #[cfg(Py_3_13)] + pub fn PyList_Clear(list: *mut PyObject) -> c_int; #[cfg_attr(PyPy, link_name = "PyPyList_Sort")] pub fn PyList_Sort(arg1: *mut PyObject) -> c_int; #[cfg_attr(PyPy, link_name = "PyPyList_Reverse")] diff --git a/src/lib.rs b/src/lib.rs index 25c88143609..9e6146eb38b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ #![warn(missing_docs)] #![cfg_attr( feature = "nightly", - feature(auto_traits, negative_impls, try_trait_v2) + feature(auto_traits, negative_impls, try_trait_v2, min_specialization) )] #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))] // Deny some lints in doctests. @@ -345,6 +345,7 @@ pub use crate::marker::Python; pub use crate::pycell::{PyRef, PyRefMut}; pub use crate::pyclass::PyClass; pub use crate::pyclass_init::PyClassInitializer; +pub use crate::try_extend::TryExtend; pub use crate::type_object::{PyTypeCheck, PyTypeInfo}; pub use crate::types::PyAny; pub use crate::version::PythonVersionInfo; @@ -449,6 +450,7 @@ pub mod pycell; pub mod pyclass; pub mod pyclass_init; +pub mod try_extend; pub mod type_object; pub mod types; mod version; diff --git a/src/prelude.rs b/src/prelude.rs index 54f5a9f6beb..759f4dc2f25 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -16,6 +16,7 @@ pub use crate::instance::{Borrowed, Bound, Py, PyObject}; pub use crate::marker::Python; pub use crate::pycell::{PyRef, PyRefMut}; pub use crate::pyclass_init::PyClassInitializer; +pub use crate::try_extend::TryExtend; pub use crate::types::{PyAny, PyModule}; #[cfg(feature = "macros")] diff --git a/src/try_extend.rs b/src/try_extend.rs new file mode 100644 index 00000000000..e4952576859 --- /dev/null +++ b/src/try_extend.rs @@ -0,0 +1,23 @@ +//! A trait for extending a collection with elements from an iterator, returning an error if the operation fails. +use crate::PyResult; + +/// A trait for extending a collection with elements from an iterator, returning an error if the operation fails. +/// This trait is similar to the standard library's `Extend` trait, but it returns a `PyResult` instead of panicking. +pub trait TryExtend +where + I: IntoIterator, +{ + /// Extends a collection with elements from an iterator, returning an error if the operation fails. + fn try_extend(&mut self, iter: I) -> PyResult<()>; +} + +impl TryExtend for T +where + I: IntoIterator, + T: Extend, +{ + fn try_extend(&mut self, iter: I) -> PyResult<()> { + self.extend(iter); + Ok(()) + } +} diff --git a/src/types/dict.rs b/src/types/dict.rs index 129f32dc9e1..905d31b39e6 100644 --- a/src/types/dict.rs +++ b/src/types/dict.rs @@ -4,7 +4,7 @@ use crate::ffi_ptr_ext::FfiPtrExt; use crate::instance::{Borrowed, Bound}; use crate::py_result_ext::PyResultExt; use crate::types::{PyAny, PyAnyMethods, PyList, PyMapping}; -use crate::{ffi, BoundObject, IntoPyObject, Python}; +use crate::{ffi, BoundObject, IntoPyObject, Python, TryExtend}; /// Represents a Python `dict`. /// @@ -408,6 +408,98 @@ impl<'py> PyDictMethods<'py> for Bound<'py, PyDict> { } } +impl<'py, I> TryExtend, Bound<'py, PyAny>)> for Bound<'_, PyDict> +where + I: IntoIterator, Bound<'py, PyAny>)>, +{ + #[cfg(not(feature = "nightly"))] + fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter() + .try_for_each(|(key, value)| self.set_item(key, value)) + } + + #[cfg(feature = "nightly")] + default fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter() + .try_for_each(|(key, value)| self.set_item(key, value)) + } +} + +impl<'py, I> TryExtend, Bound<'py, PyAny>)>> for Bound<'_, PyDict> +where + I: IntoIterator, Bound<'py, PyAny>)>>, +{ + #[cfg(not(feature = "nightly"))] + fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| { + let (key, value) = item?; + self.set_item(key, value) + }) + } + + #[cfg(feature = "nightly")] + default fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| { + let (key, value) = item?; + self.set_item(key, value) + }) + } +} + +impl<'py, I> TryExtend> for Bound<'_, PyDict> +where + I: IntoIterator>, +{ + #[cfg(not(feature = "nightly"))] + fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| { + let (key, value): (Bound<'py, PyAny>, Bound<'py, PyAny>) = item.extract()?; + self.set_item(key, value) + }) + } + + #[cfg(feature = "nightly")] + default fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| { + let (key, value): (Bound<'py, PyAny>, Bound<'py, PyAny>) = item.extract()?; + self.set_item(key, value) + }) + } +} + +#[cfg(feature = "nightly")] +impl<'py> TryExtend, (Bound<'py, PyAny>, Bound<'py, PyAny>)> + for Bound<'_, PyDict> +{ + #[cfg(feature = "nightly")] + fn try_extend(&mut self, iter: Bound<'py, PyDict>) -> PyResult<()> { + err::error_on_minusone(iter.py(), unsafe { + ffi::PyDict_Merge(self.as_ptr(), iter.as_ptr(), 1) + }) + } +} + +macro_rules! impl_try_extend_specialization( + ($i:ty, $t:ty) => { + #[cfg(feature = "nightly")] + impl<'py> TryExtend<$i, $t> for Bound<'_, PyDict> { + fn try_extend(&mut self, iter: $i) -> PyResult<()> { + err::error_on_minusone(iter.py(), unsafe { + ffi::PyDict_MergeFromSeq2(self.as_ptr(), iter.as_ptr(), 1) + }) + } + } + } +); + +impl_try_extend_specialization!( + Bound<'py, crate::types::PyIterator>, + PyResult> +); +impl_try_extend_specialization!(Bound<'py, crate::types::PyList>, Bound<'py, PyAny>); +impl_try_extend_specialization!(Bound<'py, crate::types::PySet>, Bound<'py, PyAny>); +impl_try_extend_specialization!(Bound<'py, crate::types::PyTuple>, Bound<'py, PyAny>); + impl<'a, 'py> Borrowed<'a, 'py, PyDict> { /// Iterates over the contents of this dictionary without incrementing reference counts. /// @@ -1652,4 +1744,42 @@ mod tests { .is_err()); }); } + + #[test] + fn test_dict_extend() { + Python::with_gil::<_, PyResult<()>>(|py| { + let mut dict = PyDict::new(py); + + let vec = vec![( + Bound::into_any(1.into_pyobject(py)?), + Bound::into_any(1.into_pyobject(py)?), + )]; + dict.try_extend(vec)?; + + let slice = [( + Bound::into_any(2.into_pyobject(py)?), + Bound::into_any(2.into_pyobject(py)?), + )]; + dict.try_extend(slice)?; + + let other_dict = [(3, 3)].into_py_dict(py)?; + dict.try_extend(other_dict)?; + + let list = PyList::new(py, [(4, 4)])?; + dict.try_extend(list)?; + + let tuple = PyTuple::new(py, [(5, 5)])?; + dict.try_extend(tuple)?; + + assert_eq!(dict.len(), 5); + assert!(dict.iter().all(|(k, v)| { + let k = k.extract::().unwrap(); + let v = v.extract::().unwrap(); + k == v + })); + + Ok(()) + }) + .unwrap(); + } } diff --git a/src/types/list.rs b/src/types/list.rs index f00c194739f..25152aeaa4b 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -4,7 +4,10 @@ use crate::err::{self, PyResult}; use crate::ffi::{self, Py_ssize_t}; use crate::ffi_ptr_ext::FfiPtrExt; use crate::internal_tricks::get_ssize_index; -use crate::types::{PySequence, PyTuple}; +use crate::try_extend::TryExtend; +use crate::types::{PyDict, PySequence, PyTuple}; +#[cfg(feature = "nightly")] +use crate::types::{PyIterator, PySet}; use crate::{Borrowed, Bound, BoundObject, IntoPyObject, PyAny, PyErr, PyObject, Python}; use crate::types::any::PyAnyMethods; @@ -478,6 +481,58 @@ impl<'py> PyListMethods<'py> for Bound<'py, PyList> { } } +impl<'py, I> TryExtend>> for Bound<'_, PyList> +where + I: IntoIterator>>, +{ + #[cfg(not(feature = "nightly"))] + fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| self.append(item?)) + } + + #[cfg(feature = "nightly")] + default fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| self.append(item?)) + } +} + +impl<'py, I> TryExtend> for Bound<'_, PyList> +where + I: IntoIterator>, +{ + #[cfg(not(feature = "nightly"))] + fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| self.append(item)) + } + + #[cfg(feature = "nightly")] + default fn try_extend(&mut self, iter: I) -> PyResult<()> { + iter.into_iter().try_for_each(|item| self.append(item)) + } +} + +macro_rules! impl_try_extend_specialization( + ($i:ty, $t:ty) => { + impl<'py> TryExtend<$i, $t> for Bound<'_, PyList> { + fn try_extend(&mut self, iter: $i) -> PyResult<()> { + err::error_on_minusone(self.py(), unsafe { + ffi::compat::PyList_Extend(self.as_ptr(), iter.as_ptr()) + }) + } + } + } +); + +impl_try_extend_specialization!(Bound<'py, PyDict>, (Bound<'py, PyAny>, Bound<'py, PyAny>)); +#[cfg(feature = "nightly")] +impl_try_extend_specialization!(Bound<'py, PyIterator>, PyResult>); +#[cfg(feature = "nightly")] +impl_try_extend_specialization!(Bound<'py, PyList>, Bound<'py, PyAny>); +#[cfg(feature = "nightly")] +impl_try_extend_specialization!(Bound<'py, PySet>, Bound<'py, PyAny>); +#[cfg(feature = "nightly")] +impl_try_extend_specialization!(Bound<'py, PyTuple>, Bound<'py, PyAny>); + /// Used by `PyList::iter()`. pub struct BoundListIterator<'py> { list: Bound<'py, PyList>, @@ -573,8 +628,8 @@ mod tests { use crate::types::any::PyAnyMethods; use crate::types::list::PyListMethods; use crate::types::sequence::PySequenceMethods; - use crate::types::{PyList, PyTuple}; - use crate::{ffi, IntoPyObject, Python}; + use crate::types::{IntoPyDict, PyList, PyTuple}; + use crate::{ffi, Bound, IntoPyObject, PyResult, Python, TryExtend}; #[test] fn test_new() { @@ -1102,4 +1157,30 @@ mod tests { assert!(tuple.eq(tuple_expected).unwrap()); }) } + + #[test] + fn test_list_extend() { + Python::with_gil::<_, PyResult<()>>(|py| { + let mut list = PyList::empty(py); + + let vec = vec![Bound::into_any(1.into_pyobject(py)?)]; + list.try_extend(vec)?; + + let slice = [Bound::into_any(2.into_pyobject(py)?)]; + list.try_extend(slice)?; + + let dict = [(3, 3)].into_py_dict(py)?; + list.try_extend(dict)?; + + let other_list = PyList::new(py, [4])?; + list.try_extend(other_list)?; + + let tuple = PyTuple::new(py, [5])?; + list.try_extend(tuple)?; + + assert_eq!(list.len(), 5); + Ok(()) + }) + .unwrap(); + } }