Improve PyMapping related function

pydantic_core/core_schema.py

@@ -1173,7 +1173,7 @@ def multi_host_url_schema(
     'string_too_long',
     'string_pattern_mismatch',
     'dict_type',
-    'dict_from_mapping',
+    'mapping_type',
     'list_type',
     'tuple_type',
     'set_type',

src/errors/types.rs

@@ -151,9 +151,9 @@ pub enum ErrorType {
     // dict errors
     #[strum(message = "Input should be a valid dictionary")]
     DictType,
-    #[strum(message = "Unable to convert mapping to a dictionary, error: {error}")]
-    DictFromMapping {
-        error: String,
+    #[strum(message = "Input should be a valid mapping, error: {error}")]
+    MappingType {
+        error: Cow<'static, str>,
     },
     // ---------------------
     // list errors
@@ -448,7 +448,7 @@ impl ErrorType {
             Self::StringTooShort { .. } => extract_context!(StringTooShort, ctx, min_length: usize),
             Self::StringTooLong { .. } => extract_context!(StringTooLong, ctx, max_length: usize),
             Self::StringPatternMismatch { .. } => extract_context!(StringPatternMismatch, ctx, pattern: String),
-            Self::DictFromMapping { .. } => extract_context!(DictFromMapping, ctx, error: String),
+            Self::MappingType { .. } => extract_context!(Cow::Owned, MappingType, ctx, error: String),
             Self::BytesTooShort { .. } => extract_context!(BytesTooShort, ctx, min_length: usize),
             Self::BytesTooLong { .. } => extract_context!(BytesTooLong, ctx, max_length: usize),
             Self::ValueError { .. } => extract_context!(ValueError, ctx, error: String),
@@ -540,7 +540,7 @@ impl ErrorType {
             Self::StringTooShort { min_length } => to_string_render!(self, min_length),
             Self::StringTooLong { max_length } => to_string_render!(self, max_length),
             Self::StringPatternMismatch { pattern } => render!(self, pattern),
-            Self::DictFromMapping { error } => render!(self, error),
+            Self::MappingType { error } => render!(self, error),
             Self::BytesTooShort { min_length } => to_string_render!(self, min_length),
             Self::BytesTooLong { max_length } => to_string_render!(self, max_length),
             Self::ValueError { error } => render!(self, error),
@@ -593,7 +593,7 @@ impl ErrorType {
             Self::StringTooShort { min_length } => py_dict!(py, min_length),
             Self::StringTooLong { max_length } => py_dict!(py, max_length),
             Self::StringPatternMismatch { pattern } => py_dict!(py, pattern),
-            Self::DictFromMapping { error } => py_dict!(py, error),
+            Self::MappingType { error } => py_dict!(py, error),
             Self::BytesTooShort { min_length } => py_dict!(py, min_length),
             Self::BytesTooLong { max_length } => py_dict!(py, max_length),
             Self::ValueError { error } => py_dict!(py, error),

src/input/input_python.rs

@@ -1,18 +1,17 @@
 use std::borrow::Cow;
 use std::str::from_utf8;
 
-use pyo3::exceptions::PyAttributeError;
 use pyo3::once_cell::GILOnceCell;
 use pyo3::prelude::*;
 use pyo3::types::{
     PyBool, PyByteArray, PyBytes, PyDate, PyDateTime, PyDelta, PyDict, PyFrozenSet, PyIterator, PyList, PyMapping,
-    PySequence, PySet, PyString, PyTime, PyTuple, PyType,
+    PySet, PyString, PyTime, PyTuple, PyType,
 };
 #[cfg(not(PyPy))]
 use pyo3::types::{PyDictItems, PyDictKeys, PyDictValues};
 use pyo3::{ffi, intern, AsPyPointer, PyTypeInfo};
 
-use crate::errors::{py_err_string, ErrorType, InputValue, LocItem, ValError, ValLineError, ValResult};
+use crate::errors::{ErrorType, InputValue, LocItem, ValError, ValLineError, ValResult};
 use crate::{PyMultiHostUrl, PyUrl};
 
 use super::datetime::{
@@ -329,8 +328,8 @@ impl<'a> Input<'a> for PyAny {
     fn lax_dict(&'a self) -> ValResult<GenericMapping<'a>> {
         if let Ok(dict) = self.cast_as::<PyDict>() {
             Ok(dict.into())
-        } else if let Some(generic_mapping) = mapping_as_dict(self) {
-            generic_mapping
+        } else if let Ok(mapping) = self.cast_as::<PyMapping>() {
+            Ok(mapping.into())
         } else {
             Err(ValError::new(ErrorType::DictType, self))
         }
@@ -342,9 +341,8 @@ impl<'a> Input<'a> for PyAny {
             if let Ok(dict) = self.cast_as::<PyDict>() {
                 return Ok(dict.into());
             } else if !strict {
-                // we can't do this in one set of if/else because we need to check from_mapping before doing this
-                if let Some(generic_mapping) = mapping_as_dict(self) {
-                    return generic_mapping;
+                if let Ok(mapping) = self.cast_as::<PyMapping>() {
+                    return Ok(mapping.into());
                 }
             }
 
@@ -643,47 +641,6 @@ impl<'a> Input<'a> for PyAny {
     }
 }
 
-/// return None if obj is not a mapping (cast_as::<PyMapping> fails or mapping.items returns an AttributeError)
-/// otherwise try to covert the mapping to a dict and return an Some(error) if it fails
-fn mapping_as_dict(obj: &PyAny) -> Option<ValResult<GenericMapping>> {
-    let mapping: &PyMapping = match obj.cast_as() {
-        Ok(mapping) => mapping,
-        Err(_) => return None,
-    };
-    // see https://github.com/PyO3/pyo3/issues/2072 - the cast_as::<PyMapping> is not entirely accurate
-    // and returns some things which are definitely not mappings (e.g. str) as mapping,
-    // hence we also require that the object as `items` to consider it a mapping
-    let result_dict = match mapping.items() {
-        Ok(seq) => mapping_seq_as_dict(seq),
-        Err(err) => {
-            if matches!(err.get_type(obj.py()).is_subclass_of::<PyAttributeError>(), Ok(true)) {
-                return None;
-            } else {
-                Err(err)
-            }
-        }
-    };
-    match result_dict {
-        Ok(dict) => Some(Ok(dict.into())),
-        Err(err) => Some(Err(ValError::new(
-            ErrorType::DictFromMapping {
-                error: py_err_string(obj.py(), err),
-            },
-            obj,
-        ))),
-    }
-}
-
-// creating a temporary dict is slow, we could perhaps use an indexmap instead
-fn mapping_seq_as_dict(seq: &PySequence) -> PyResult<&PyDict> {
-    let dict = PyDict::new(seq.py());
-    for r in seq.iter()? {
-        let (key, value): (&PyAny, &PyAny) = r?.extract()?;
-        dict.set_item(key, value)?;
-    }
-    Ok(dict)
-}
-
 /// Best effort check of whether it's likely to make sense to inspect obj for attributes and iterate over it
 /// with `obj.dir()`
 fn from_attributes_applicable(obj: &PyAny) -> bool {

src/input/mod.rs

@@ -14,8 +14,9 @@ pub use datetime::{EitherDate, EitherDateTime, EitherTime, EitherTimedelta};
 pub use input_abstract::Input;
 pub use parse_json::{JsonInput, JsonObject, JsonType};
 pub use return_enums::{
-    py_string_str, EitherBytes, EitherString, GenericArguments, GenericCollection, GenericIterator, GenericMapping,
-    JsonArgs, PyArgs,
+    py_string_str, AttributesGenericIterator, DictGenericIterator, EitherBytes, EitherString, GenericArguments,
+    GenericCollection, GenericIterator, GenericMapping, JsonArgs, JsonObjectGenericIterator, MappingGenericIterator,
+    PyArgs,
 };
 
 pub fn repr_string(v: &PyAny) -> PyResult<String> {

src/input/return_enums.rs

@@ -1,7 +1,16 @@
 use std::borrow::Cow;
 
+use pyo3::intern;
 use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyDict, PyFrozenSet, PyIterator, PyList, PySet, PyString, PyTuple};
+use pyo3::types::iter::PyDictIterator;
+use pyo3::types::{PyBytes, PyDict, PyFrozenSet, PyIterator, PyList, PyMapping, PySet, PyString, PyTuple};
+
+#[cfg(not(PyPy))]
+use pyo3::types::PyFunction;
+#[cfg(not(PyPy))]
+use pyo3::PyTypeInfo;
+
+use indexmap::map::Iter;
 
 use crate::errors::{py_err_string, ErrorType, InputValue, ValError, ValLineError, ValResult};
 use crate::recursion_guard::RecursionGuard;
@@ -233,14 +242,187 @@ impl<'a> GenericCollection<'a> {
 #[cfg_attr(debug_assertions, derive(Debug))]
 pub enum GenericMapping<'a> {
     PyDict(&'a PyDict),
+    PyMapping(&'a PyMapping),
     PyGetAttr(&'a PyAny),
     JsonObject(&'a JsonObject),
 }
 
 derive_from!(GenericMapping, PyDict, PyDict);
+derive_from!(GenericMapping, PyMapping, PyMapping);
 derive_from!(GenericMapping, PyGetAttr, PyAny);
 derive_from!(GenericMapping, JsonObject, JsonObject);
 
+pub struct DictGenericIterator<'py> {
+    dict_iter: PyDictIterator<'py>,
+}
+
+impl<'py> DictGenericIterator<'py> {
+    pub fn new(dict: &'py PyDict) -> ValResult<'py, Self> {
+        Ok(Self { dict_iter: dict.iter() })
+    }
+}
+
+impl<'py> Iterator for DictGenericIterator<'py> {
+    type Item = ValResult<'py, (&'py PyAny, &'py PyAny)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.dict_iter.next().map(Ok)
+    }
+    // size_hint is omitted as it isn't needed
+}
+
+pub struct MappingGenericIterator<'py> {
+    input: &'py PyAny,
+    iter: &'py PyIterator,
+}
+
+fn mapping_err<'py>(err: PyErr, py: Python<'py>, input: &'py impl Input<'py>) -> ValError<'py> {
+    ValError::new(
+        ErrorType::MappingType {
+            error: py_err_string(py, err).into(),
+        },
+        input,
+    )
+}
+
+impl<'py> MappingGenericIterator<'py> {
+    pub fn new(mapping: &'py PyMapping) -> ValResult<'py, Self> {
+        let py = mapping.py();
+        let input: &PyAny = mapping;
+        let iter = mapping
+            .items()
+            .map_err(|e| mapping_err(e, py, input))?
+            .iter()
+            .map_err(|e| mapping_err(e, py, input))?;
+        Ok(Self { iter, input })
+    }
+}
+
+const MAPPING_TUPLE_ERROR: &str = "Mapping items must be tuples of (key, value) pairs";
+
+impl<'py> Iterator for MappingGenericIterator<'py> {
+    type Item = ValResult<'py, (&'py PyAny, &'py PyAny)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let item = match self.iter.next() {
+            Some(Err(e)) => return Some(Err(mapping_err(e, self.iter.py(), self.input))),
+            Some(Ok(item)) => item,
+            None => return None,
+        };
+        let tuple: &PyTuple = match item.cast_as() {
+            Ok(tuple) => tuple,
+            Err(_) => {
+                return Some(Err(ValError::new(
+                    ErrorType::MappingType {
+                        error: MAPPING_TUPLE_ERROR.into(),
+                    },
+                    self.input,
+                )))
+            }
+        };
+        if tuple.len() != 2 {
+            return Some(Err(ValError::new(
+                ErrorType::MappingType {
+                    error: MAPPING_TUPLE_ERROR.into(),
+                },
+                self.input,
+            )));
+        };
+        #[cfg(PyPy)]
+        let key = tuple.get_item(0).unwrap();
+        #[cfg(PyPy)]
+        let value = tuple.get_item(1).unwrap();
+        #[cfg(not(PyPy))]
+        let key = unsafe { tuple.get_item_unchecked(0) };
+        #[cfg(not(PyPy))]
+        let value = unsafe { tuple.get_item_unchecked(1) };
+        Some(Ok((key, value)))
+    }
+    // size_hint is omitted as it isn't needed
+}
+
+pub struct AttributesGenericIterator<'py> {
+    object: &'py PyAny,
+    attributes: &'py PyList,
+    index: usize,
+}
+
+impl<'py> AttributesGenericIterator<'py> {
+    pub fn new(py_any: &'py PyAny) -> ValResult<'py, Self> {
+        Ok(Self {
+            object: py_any,
+            attributes: py_any.dir(),
+            index: 0,
+        })
+    }
+}
+
+impl<'py> Iterator for AttributesGenericIterator<'py> {
+    type Item = ValResult<'py, (&'py PyAny, &'py PyAny)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // loop until we find an attribute who's name does not start with underscore,
+        // or we get to the end of the list of attributes
+        while self.index < self.attributes.len() {
+            #[cfg(PyPy)]
+            let name: &PyAny = self.attributes.get_item(self.index).unwrap();
+            #[cfg(not(PyPy))]
+            let name: &PyAny = unsafe { self.attributes.get_item_unchecked(self.index) };
+            self.index += 1;
+            // from benchmarks this is 14x faster than using the python `startswith` method
+            let name_cow = match name.cast_as::<PyString>() {
+                Ok(name) => name.to_string_lossy(),
+                Err(e) => return Some(Err(e.into())),
+            };
+            if !name_cow.as_ref().starts_with('_') {
+                // getattr is most likely to fail due to an exception in a @property, skip
+                if let Ok(attr) = self.object.getattr(name_cow.as_ref()) {
+                    // we don't want bound methods to be included, is there a better way to check?
+                    // ref https://stackoverflow.com/a/18955425/949890
+                    let is_bound = matches!(attr.hasattr(intern!(attr.py(), "__self__")), Ok(true));
+                    // the PyFunction::is_type_of(attr) catches `staticmethod`, but also any other function,
+                    // I think that's better than including static methods in the yielded attributes,
+                    // if someone really wants fields, they can use an explicit field, or a function to modify input
+                    #[cfg(not(PyPy))]
+                    if !is_bound && !PyFunction::is_type_of(attr) {
+                        return Some(Ok((name, attr)));
+                    }
+                    // MASSIVE HACK! PyFunction doesn't exist for PyPy,
+                    // is_instance_of::<PyFunction> crashes with a null pointer, hence this hack, see
+                    // https://github.com/pydantic/pydantic-core/pull/161#discussion_r917257635
+                    #[cfg(PyPy)]
+                    if !is_bound && attr.get_type().to_string() != "<class 'function'>" {
+                        return Some(Ok((name, attr)));
+                    }
+                }
+            }
+        }
+        None
+    }
+    // size_hint is omitted as it isn't needed
+}
+
+pub struct JsonObjectGenericIterator<'py> {
+    object_iter: Iter<'py, String, JsonInput>,
+}
+
+impl<'py> JsonObjectGenericIterator<'py> {
+    pub fn new(json_object: &'py JsonObject) -> ValResult<'py, Self> {
+        Ok(Self {
+            object_iter: json_object.iter(),
+        })
+    }
+}
+
+impl<'py> Iterator for JsonObjectGenericIterator<'py> {
+    type Item = ValResult<'py, (&'py String, &'py JsonInput)>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.object_iter.next().map(Ok)
+    }
+    // size_hint is omitted as it isn't needed
+}
+
 #[derive(Debug, Clone)]
 pub enum GenericIterator {
     PyIterator(GenericPyIterator),