Add TryExtend trait

newsfragments/4667.added.md

@@ -0,0 +1 @@
+Add `PyList_Extend` & `PyList_Clear` to pyo3-ffi

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())
+    }
+);

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")]

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;

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")]

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<I, T>
+where
+    I: IntoIterator<Item = T>,
+{
+    /// Extends a collection with elements from an iterator, returning an error if the operation fails.
+    fn try_extend(&mut self, iter: I) -> PyResult<()>;
+}
+
+impl<I, A, T> TryExtend<I, A> for T
+where
+    I: IntoIterator<Item = A>,
+    T: Extend<A>,
+{
+    fn try_extend(&mut self, iter: I) -> PyResult<()> {
+        self.extend(iter);
+        Ok(())
+    }
+}

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<I, (Bound<'py, PyAny>, Bound<'py, PyAny>)> for Bound<'_, PyDict>
+where
+    I: IntoIterator<Item = (Bound<'py, PyAny>, 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<I, PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)>> for Bound<'_, PyDict>
+where
+    I: IntoIterator<Item = PyResult<(Bound<'py, PyAny>, 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<I, Bound<'py, PyAny>> for Bound<'_, PyDict>
+where
+    I: IntoIterator<Item = 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): (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, PyDict>, (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<Bound<'py, PyAny>>
+);
+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::<i32>().unwrap();
+                let v = v.extract::<i32>().unwrap();
+                k == v
+            }));
+
+            Ok(())
+        })
+        .unwrap();
+    }
 }

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<I, PyResult<Bound<'py, PyAny>>> for Bound<'_, PyList>
+where
+    I: IntoIterator<Item = PyResult<Bound<'py, PyAny>>>,
+{
+    #[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<I, Bound<'py, PyAny>> for Bound<'_, PyList>
+where
+    I: IntoIterator<Item = Bound<'py, PyAny>>,
+{
+    #[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<Bound<'py, PyAny>>);
+#[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();
+    }
 }