rework to handle sets

Author: adriangb

benches/main.rs

@@ -5,7 +5,7 @@ extern crate test;
 use test::{black_box, Bencher};
 
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyString};
+use pyo3::types::{PyDict, PySet, PyString};
 
 use _pydantic_core::SchemaValidator;
 
@@ -696,3 +696,58 @@ class Foo(Enum):
         }
     })
 }
+
+const COLLECTION_SIZE: usize = 100_000;
+
+#[bench]
+fn constructing_pyset_from_vec_without_capacity(bench: &mut Bencher) {
+    Python::with_gil(|py| {
+        let input: Vec<PyObject> = (0..COLLECTION_SIZE).map(|v| v.to_object(py)).collect();
+
+        bench.iter(|| {
+            black_box({
+                let mut output = Vec::new();
+                for x in &input {
+                    output.push(x);
+                }
+                let set = PySet::new(py, output.iter()).unwrap();
+                set
+            })
+        })
+    })
+}
+
+#[bench]
+fn constructing_pyset_from_vec_with_capacity(bench: &mut Bencher) {
+    Python::with_gil(|py| {
+        let input: Vec<PyObject> = (0..COLLECTION_SIZE).map(|v| v.to_object(py)).collect();
+
+        bench.iter(|| {
+            black_box({
+                let mut output = Vec::with_capacity(COLLECTION_SIZE);
+                for x in &input {
+                    output.push(x);
+                }
+                let set = PySet::new(py, output.iter()).unwrap();
+                set
+            })
+        })
+    })
+}
+
+#[bench]
+fn constructing_pyset_from_vec_directly(bench: &mut Bencher) {
+    Python::with_gil(|py| {
+        let input: Vec<PyObject> = (0..COLLECTION_SIZE).map(|v| v.to_object(py)).collect();
+
+        bench.iter(|| {
+            black_box({
+                let output = PySet::new(py, &Vec::<i64>::new()).unwrap();
+                for x in &input {
+                    output.add(x).unwrap();
+                }
+                output
+            })
+        })
+    })
+}

src/input/any_iterable.rs

@@ -1,7 +1,14 @@
 use super::parse_json::{JsonInput, JsonObject};
-use pyo3::types::{PyDict, PyFrozenSet, PyIterator, PyList, PyMapping, PySequence, PySet, PyTuple};
+use pyo3::{
+    exceptions::PyTypeError,
+    types::{
+        PyByteArray, PyBytes, PyDict, PyFrozenSet, PyIterator, PyList, PyMapping, PySequence, PySet, PyString, PyTuple,
+    },
+    PyAny, PyResult, Python, ToPyObject,
+};
 
-pub enum AnyIterable<'a> {
+#[derive(Debug)]
+pub enum GenericIterable<'a> {
     List(&'a PyList),
     Tuple(&'a PyTuple),
     Set(&'a PySet),
@@ -13,8 +20,100 @@ pub enum AnyIterable<'a> {
     DictValues(&'a PyIterator),
     DictItems(&'a PyIterator),
     Mapping(&'a PyMapping),
+    String(&'a PyString),
+    Bytes(&'a PyBytes),
+    PyByteArray(&'a PyByteArray),
     Sequence(&'a PySequence),
     Iterator(&'a PyIterator),
     JsonArray(&'a [JsonInput]),
     JsonObject(&'a JsonObject),
 }
+
+type PyMappingItems<'a> = (&'a PyAny, &'a PyAny);
+
+#[inline(always)]
+fn extract_items(item: PyResult<&PyAny>) -> PyResult<PyMappingItems<'_>> {
+    match item {
+        Ok(v) => v.extract::<PyMappingItems>(),
+        Err(e) => Err(e),
+    }
+}
+
+impl<'a, 'py: 'a> GenericIterable<'a> {
+    pub fn into_sequence_iterator(
+        self,
+        py: Python<'py>,
+    ) -> PyResult<Box<dyn Iterator<Item = PyResult<&'a PyAny>> + 'a>> {
+        match self {
+            GenericIterable::List(iter) => Ok(Box::new(iter.iter().map(Ok))),
+            GenericIterable::Tuple(iter) => Ok(Box::new(iter.iter().map(Ok))),
+            GenericIterable::Set(iter) => Ok(Box::new(iter.iter().map(Ok))),
+            GenericIterable::FrozenSet(iter) => Ok(Box::new(iter.iter().map(Ok))),
+            // Note that this iterates over only the keys, just like doing iter({}) in Python
+            GenericIterable::Dict(iter) => Ok(Box::new(iter.iter().map(|(k, _)| Ok(k)))),
+            GenericIterable::DictKeys(iter) => Ok(Box::new(iter)),
+            GenericIterable::DictValues(iter) => Ok(Box::new(iter)),
+            GenericIterable::DictItems(iter) => Ok(Box::new(iter)),
+            // Note that this iterates over only the keys, just like doing iter({}) in Python
+            GenericIterable::Mapping(iter) => Ok(Box::new(iter.keys()?.iter()?)),
+            GenericIterable::String(iter) => Ok(Box::new(iter.iter()?)),
+            GenericIterable::Bytes(iter) => Ok(Box::new(iter.iter()?)),
+            GenericIterable::PyByteArray(iter) => Ok(Box::new(iter.iter()?)),
+            GenericIterable::Sequence(iter) => Ok(Box::new(iter.iter()?)),
+            GenericIterable::Iterator(iter) => Ok(Box::new(iter)),
+            GenericIterable::JsonArray(iter) => Ok(Box::new(iter.iter().map(move |v| {
+                let v = v.to_object(py);
+                Ok(v.into_ref(py))
+            }))),
+            // Note that this iterates over only the keys, just like doing iter({}) in Python, just for consistency
+            GenericIterable::JsonObject(iter) => Ok(Box::new(
+                iter.iter().map(move |(k, _)| Ok(k.to_object(py).into_ref(py))),
+            )),
+        }
+    }
+
+    pub fn into_mapping_items_iterator(
+        self,
+        py: Python<'py>,
+    ) -> PyResult<Box<dyn Iterator<Item = PyResult<PyMappingItems<'a>>> + 'a>> {
+        match self {
+            GenericIterable::List(iter) => Ok(Box::new(iter.iter().map(|v| extract_items(Ok(v))))),
+            GenericIterable::Tuple(iter) => Ok(Box::new(iter.iter().map(|v| extract_items(Ok(v))))),
+            GenericIterable::Set(iter) => Ok(Box::new(iter.iter().map(|v| extract_items(Ok(v))))),
+            GenericIterable::FrozenSet(iter) => Ok(Box::new(iter.iter().map(|v| extract_items(Ok(v))))),
+            // Note that we iterate over (key, value), unlike doing iter({}) in Python
+            GenericIterable::Dict(iter) => Ok(Box::new(iter.iter().map(Ok))),
+            // Keys or values can be tuples
+            GenericIterable::DictKeys(iter) => Ok(Box::new(iter.map(extract_items))),
+            GenericIterable::DictValues(iter) => Ok(Box::new(iter.map(extract_items))),
+            GenericIterable::DictItems(iter) => Ok(Box::new(iter.map(extract_items))),
+            // Note that we iterate over (key, value), unlike doing iter({}) in Python
+            GenericIterable::Mapping(iter) => Ok(Box::new(iter.items()?.iter()?.map(extract_items))),
+            // In Python if you do dict("foobar") you get "dictionary update sequence element #0 has length 1; 2 is required"
+            // This is similar but arguably a better error message
+            GenericIterable::String(_) => Err(PyTypeError::new_err(
+                "Expected an iterable of (key, value) pairs, got a string",
+            )),
+            GenericIterable::Bytes(_) => Err(PyTypeError::new_err(
+                "Expected an iterable of (key, value) pairs, got a bytes",
+            )),
+            GenericIterable::PyByteArray(_) => Err(PyTypeError::new_err(
+                "Expected an iterable of (key, value) pairs, got a bytearray",
+            )),
+            // Obviously these may be things that are not convertible to a tuple of (Hashable, Any)
+            // Python fails with a similar error message to above, ours will be slightly different (PyO3 will fail to extract) but similar enough
+            GenericIterable::Sequence(iter) => Ok(Box::new(iter.iter()?.map(extract_items))),
+            GenericIterable::Iterator(iter) => Ok(Box::new(iter.iter()?.map(extract_items))),
+            GenericIterable::JsonArray(iter) => Ok(Box::new(
+                iter.iter()
+                    .map(move |v| extract_items(Ok(v.to_object(py).into_ref(py)))),
+            )),
+            // Note that we iterate over (key, value), unlike doing iter({}) in Python
+            GenericIterable::JsonObject(iter) => Ok(Box::new(iter.iter().map(move |(k, v)| {
+                let k = PyString::new(py, k).as_ref();
+                let v = v.to_object(py).into_ref(py);
+                Ok((k, v))
+            }))),
+        }
+    }
+}

src/input/input_abstract.rs

@@ -6,7 +6,7 @@ use pyo3::{intern, prelude::*};
 use crate::errors::{InputValue, LocItem, ValResult};
 use crate::{PyMultiHostUrl, PyUrl};
 
-use super::any_iterable::AnyIterable;
+use super::any_iterable::GenericIterable;
 use super::datetime::{EitherDate, EitherDateTime, EitherTime, EitherTimedelta};
 use super::return_enums::{EitherBytes, EitherString};
 use super::{GenericArguments, GenericCollection, GenericIterator, GenericMapping, JsonInput};
@@ -167,21 +167,7 @@ pub trait Input<'a>: fmt::Debug + ToPyObject {
         self.validate_dict(strict)
     }
 
-    fn validate_list(&'a self, strict: bool, allow_any_iter: bool) -> ValResult<GenericCollection<'a>> {
-        if strict && !allow_any_iter {
-            self.strict_list()
-        } else {
-            self.lax_list(allow_any_iter)
-        }
-    }
-
-    fn strict_list(&'a self) -> ValResult<GenericCollection<'a>>;
-    #[cfg_attr(has_no_coverage, no_coverage)]
-    fn lax_list(&'a self, _allow_any_iter: bool) -> ValResult<GenericCollection<'a>> {
-        self.strict_list()
-    }
-
-    fn extract_iterable(&'a self) -> ValResult<AnyIterable<'a>>;
+    fn extract_iterable(&'a self) -> ValResult<GenericIterable<'a>>;
 
     fn validate_tuple(&'a self, strict: bool) -> ValResult<GenericCollection<'a>> {
         if strict {

src/input/input_json.rs

@@ -187,17 +187,6 @@ impl<'a> Input<'a> for JsonInput {
         self.validate_dict(false)
     }
 
-    fn validate_list(&'a self, _strict: bool, _allow_any_iter: bool) -> ValResult<GenericCollection<'a>> {
-        match self {
-            JsonInput::Array(a) => Ok(a.into()),
-            _ => Err(ValError::new(ErrorType::ListType, self)),
-        }
-    }
-    #[cfg_attr(has_no_coverage, no_coverage)]
-    fn strict_list(&'a self) -> ValResult<GenericCollection<'a>> {
-        self.validate_list(false, false)
-    }
-
     fn validate_tuple(&'a self, _strict: bool) -> ValResult<GenericCollection<'a>> {
         // just as in set's case, List has to be allowed
         match self {
@@ -234,10 +223,10 @@ impl<'a> Input<'a> for JsonInput {
         self.validate_frozenset(false)
     }
 
-    fn extract_iterable(&'a self) -> ValResult<super::any_iterable::AnyIterable<'a>> {
+    fn extract_iterable(&'a self) -> ValResult<super::any_iterable::GenericIterable<'a>> {
         match self {
-            JsonInput::Array(a) => Ok(super::any_iterable::AnyIterable::JsonArray(a)),
-            JsonInput::Object(o) => Ok(super::any_iterable::AnyIterable::JsonObject(o)),
+            JsonInput::Array(a) => Ok(super::any_iterable::GenericIterable::JsonArray(a)),
+            JsonInput::Object(o) => Ok(super::any_iterable::GenericIterable::JsonObject(o)),
             _ => Err(ValError::new(ErrorType::IterableType, self)),
         }
     }
@@ -412,16 +401,7 @@ impl<'a> Input<'a> for String {
         self.validate_dict(false)
     }
 
-    #[cfg_attr(has_no_coverage, no_coverage)]
-    fn validate_list(&'a self, _strict: bool, _allow_any_iter: bool) -> ValResult<GenericCollection<'a>> {
-        Err(ValError::new(ErrorType::ListType, self))
-    }
-    #[cfg_attr(has_no_coverage, no_coverage)]
-    fn strict_list(&'a self) -> ValResult<GenericCollection<'a>> {
-        self.validate_list(false, false)
-    }
-
-    fn extract_iterable(&'a self) -> ValResult<super::any_iterable::AnyIterable<'a>> {
+    fn extract_iterable(&'a self) -> ValResult<super::any_iterable::GenericIterable<'a>> {
         Err(ValError::new(ErrorType::IterableType, self))
     }
 

src/input/input_python.rs

@@ -11,7 +11,7 @@ use pyo3::types::{
 use pyo3::types::{PyDictItems, PyDictKeys, PyDictValues};
 use pyo3::{ffi, intern, AsPyPointer, PyTypeInfo};
 
-use super::any_iterable::AnyIterable;
+use super::any_iterable::GenericIterable;
 use crate::build_tools::safe_repr;
 use crate::errors::{ErrorType, InputValue, LocItem, ValError, ValResult};
 use crate::{ArgsKwargs, PyMultiHostUrl, PyUrl};
@@ -437,30 +437,6 @@ impl<'a> Input<'a> for PyAny {
         }
     }
 
-    fn strict_list(&'a self) -> ValResult<GenericCollection<'a>> {
-        if let Ok(list) = self.downcast::<PyList>() {
-            Ok(list.into())
-        } else {
-            Err(ValError::new(ErrorType::ListType, self))
-        }
-    }
-
-    fn lax_list(&'a self, allow_any_iter: bool) -> ValResult<GenericCollection<'a>> {
-        if let Ok(list) = self.downcast::<PyList>() {
-            Ok(list.into())
-        } else if let Ok(tuple) = self.downcast::<PyTuple>() {
-            Ok(tuple.into())
-        } else if let Some(collection) = extract_dict_iter!(self) {
-            Ok(collection)
-        } else if allow_any_iter && self.iter().is_ok() {
-            Ok(self.into())
-        } else if let Some(collection) = extract_shared_iter!(PyList, self) {
-            Ok(collection)
-        } else {
-            Err(ValError::new(ErrorType::ListType, self))
-        }
-    }
-
     fn strict_tuple(&'a self) -> ValResult<GenericCollection<'a>> {
         if let Ok(tuple) = self.downcast::<PyTuple>() {
             Ok(tuple.into())
@@ -535,37 +511,36 @@ impl<'a> Input<'a> for PyAny {
         }
     }
 
-    fn extract_iterable(&'a self) -> ValResult<super::any_iterable::AnyIterable<'a>> {
+    fn extract_iterable(&'a self) -> ValResult<super::any_iterable::GenericIterable<'a>> {
         // Handle concrete non-overlapping types first, then abstract types
         if let Ok(iterable) = self.downcast::<PyList>() {
-            Ok(AnyIterable::List(iterable))
+            Ok(GenericIterable::List(iterable))
         } else if let Ok(iterable) = self.downcast::<PyTuple>() {
-            Ok(AnyIterable::Tuple(iterable))
+            Ok(GenericIterable::Tuple(iterable))
         } else if let Ok(iterable) = self.downcast::<PySet>() {
-            Ok(AnyIterable::Set(iterable))
+            Ok(GenericIterable::Set(iterable))
         } else if let Ok(iterable) = self.downcast::<PyFrozenSet>() {
-            Ok(AnyIterable::FrozenSet(iterable))
+            Ok(GenericIterable::FrozenSet(iterable))
         } else if let Ok(iterable) = self.downcast::<PyDict>() {
-            Ok(AnyIterable::Dict(iterable))
+            Ok(GenericIterable::Dict(iterable))
         } else if let Some(iterable) = extract_dict_keys!(self.py(), self) {
-            Ok(AnyIterable::DictKeys(iterable))
+            Ok(GenericIterable::DictKeys(iterable))
         } else if let Some(iterable) = extract_dict_values!(self.py(), self) {
-            Ok(AnyIterable::DictValues(iterable))
+            Ok(GenericIterable::DictValues(iterable))
         } else if let Some(iterable) = extract_dict_items!(self.py(), self) {
-            Ok(AnyIterable::DictItems(iterable))
+            Ok(GenericIterable::DictItems(iterable))
         } else if let Ok(iterable) = self.downcast::<PyMapping>() {
-            Ok(AnyIterable::Mapping(iterable))
-        } else if let (Ok(iterable), Err(_), Err(_)) = (
-            self.downcast::<PySequence>(),
-            // Explicitly disallow strings and bytes since they are sequences
-            // but you almost never want to treat it as one
-            // This can be worked around by allowing arbitrary iterables
-            self.downcast::<PyString>(),
-            self.downcast::<PyBytes>(),
-        ) {
-            Ok(AnyIterable::Sequence(iterable))
-        } else if let Ok(iterable) = PyIterator::from_object(self.py(), self) {
-            Ok(AnyIterable::Iterator(iterable))
+            Ok(GenericIterable::Mapping(iterable))
+        } else if let Ok(iterable) = self.downcast::<PyString>() {
+            Ok(GenericIterable::String(iterable))
+        } else if let Ok(iterable) = self.downcast::<PyBytes>() {
+            Ok(GenericIterable::Bytes(iterable))
+        } else if let Ok(iterable) = self.downcast::<PyByteArray>() {
+            Ok(GenericIterable::PyByteArray(iterable))
+        } else if let Ok(iterable) = self.downcast::<PySequence>() {
+            Ok(GenericIterable::Sequence(iterable))
+        } else if let Ok(iterable) = self.iter() {
+            Ok(GenericIterable::Iterator(iterable))
         } else {
             Err(ValError::new(ErrorType::IterableType, self))
         }