fix set length validation

Author: samuelcolvin

pydantic_core/core_schema.py

@@ -1353,7 +1353,6 @@ class SetSchema(TypedDict, total=False):
     items_schema: CoreSchema
     min_length: int
     max_length: int
-    generator_max_length: int
     strict: bool
     ref: str
     metadata: Any
@@ -1365,7 +1364,6 @@ def set_schema(
     *,
     min_length: int | None = None,
     max_length: int | None = None,
-    generator_max_length: int | None = None,
     strict: bool | None = None,
     ref: str | None = None,
     metadata: Any = None,
@@ -1388,9 +1386,6 @@ def set_schema(
         items_schema: The value must be a set with items that match this schema
         min_length: The value must be a set with at least this many items
         max_length: The value must be a set with at most this many items
-        generator_max_length: At most this many items will be read from a generator before failing validation
-            This is important because generators can be infinite, and even with a `max_length` on the set,
-            an infinite generator could run forever without producing more than `max_length` distinct items.
         strict: The value must be a set with exactly this many items
         ref: optional unique identifier of the schema, used to reference the schema in other places
         metadata: Any other information you want to include with the schema, not used by pydantic-core
@@ -1401,7 +1396,6 @@ def set_schema(
         items_schema=items_schema,
         min_length=min_length,
         max_length=max_length,
-        generator_max_length=generator_max_length,
         strict=strict,
         ref=ref,
         metadata=metadata,
@@ -1414,7 +1408,6 @@ class FrozenSetSchema(TypedDict, total=False):
     items_schema: CoreSchema
     min_length: int
     max_length: int
-    generator_max_length: int
     strict: bool
     ref: str
     metadata: Any
@@ -1426,7 +1419,6 @@ def frozenset_schema(
     *,
     min_length: int | None = None,
     max_length: int | None = None,
-    generator_max_length: int | None = None,
     strict: bool | None = None,
     ref: str | None = None,
     metadata: Any = None,
@@ -1449,7 +1441,6 @@ def frozenset_schema(
         items_schema: The value must be a frozenset with items that match this schema
         min_length: The value must be a frozenset with at least this many items
         max_length: The value must be a frozenset with at most this many items
-        generator_max_length: The value must generate a frozenset with at most this many items
         strict: The value must be a frozenset with exactly this many items
         ref: optional unique identifier of the schema, used to reference the schema in other places
         metadata: Any other information you want to include with the schema, not used by pydantic-core
@@ -1460,7 +1451,6 @@ def frozenset_schema(
         items_schema=items_schema,
         min_length=min_length,
         max_length=max_length,
-        generator_max_length=generator_max_length,
         strict=strict,
         ref=ref,
         metadata=metadata,

src/input/return_enums.rs

@@ -131,6 +131,54 @@ fn validate_iter_to_vec<'a, 's>(
     }
 }
 
+#[allow(clippy::too_many_arguments)]
+fn validate_iter_to_set<'a, 's>(
+    py: Python<'a>,
+    iter: impl Iterator<Item = PyResult<&'a (impl Input<'a> + 'a)>>,
+    input: &'a (impl Input<'a> + 'a),
+    field_type: &'static str,
+    max_length: Option<usize>,
+    validator: &'s CombinedValidator,
+    extra: &Extra,
+    definitions: &'a [CombinedValidator],
+    recursion_guard: &'s mut RecursionGuard,
+) -> ValResult<'a, &'a PySet> {
+    let set = PySet::empty(py)?;
+    let mut errors: Vec<ValLineError> = Vec::new();
+    for (index, item_result) in iter.enumerate() {
+        let item = item_result.map_err(|e| any_next_error!(py, e, input, index))?;
+        match validator.validate(py, item, extra, definitions, recursion_guard) {
+            Ok(item) => {
+                set.add(item)?;
+                if let Some(max_length) = max_length {
+                    let actual_length = set.len();
+                    if actual_length > max_length {
+                        return Err(ValError::new(
+                            ErrorType::TooLong {
+                                field_type: field_type.to_string(),
+                                max_length,
+                                actual_length,
+                            },
+                            input,
+                        ));
+                    }
+                }
+            }
+            Err(ValError::LineErrors(line_errors)) => {
+                errors.extend(line_errors.into_iter().map(|err| err.with_outer_location(index.into())));
+            }
+            Err(ValError::Omit) => (),
+            Err(err) => return Err(err),
+        }
+    }
+
+    if errors.is_empty() {
+        Ok(set)
+    } else {
+        Err(ValError::LineErrors(errors))
+    }
+}
+
 fn no_validator_iter_to_vec<'a, 's>(
     py: Python<'a>,
     iter: impl Iterator<Item = &'a (impl Input<'a> + 'a)>,
@@ -174,67 +222,67 @@ impl<'a> GenericCollection<'a> {
             .generic_len()
             .unwrap_or_else(|_| max_length.unwrap_or(DEFAULT_CAPACITY));
         let max_length_check = MaxLengthCheck::new(max_length, field_type, input);
+
+        macro_rules! validate {
+            ($iter:expr) => {
+                validate_iter_to_vec(
+                    py,
+                    $iter,
+                    capacity,
+                    max_length_check,
+                    validator,
+                    extra,
+                    definitions,
+                    recursion_guard,
+                )
+            };
+        }
+
+        match self {
+            Self::List(collection) => validate!(collection.iter().map(Ok)),
+            Self::Tuple(collection) => validate!(collection.iter().map(Ok)),
+            Self::Set(collection) => validate!(collection.iter().map(Ok)),
+            Self::FrozenSet(collection) => validate!(collection.iter().map(Ok)),
+            Self::PyAny(collection) => validate!(collection.iter()?),
+            Self::JsonArray(collection) => validate!(collection.iter().map(Ok)),
+        }
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn validate_to_set<'s>(
+        &'s self,
+        py: Python<'a>,
+        input: &'a impl Input<'a>,
+        max_length: Option<usize>,
+        field_type: &'static str,
+        validator: &'s CombinedValidator,
+        extra: &Extra,
+        definitions: &'a [CombinedValidator],
+        recursion_guard: &'s mut RecursionGuard,
+    ) -> ValResult<'a, &'a PySet> {
+        macro_rules! validate_set {
+            ($iter:expr) => {
+                validate_iter_to_set(
+                    py,
+                    $iter,
+                    input,
+                    field_type,
+                    max_length,
+                    validator,
+                    extra,
+                    definitions,
+                    recursion_guard,
+                )
+            };
+        }
+
         match self {
-            Self::List(collection) => validate_iter_to_vec(
-                py,
-                collection.iter().map(Ok),
-                capacity,
-                max_length_check,
-                validator,
-                extra,
-                definitions,
-                recursion_guard,
-            ),
-            Self::Tuple(collection) => validate_iter_to_vec(
-                py,
-                collection.iter().map(Ok),
-                capacity,
-                max_length_check,
-                validator,
-                extra,
-                definitions,
-                recursion_guard,
-            ),
-            Self::Set(collection) => validate_iter_to_vec(
-                py,
-                collection.iter().map(Ok),
-                capacity,
-                max_length_check,
-                validator,
-                extra,
-                definitions,
-                recursion_guard,
-            ),
-            Self::FrozenSet(collection) => validate_iter_to_vec(
-                py,
-                collection.iter().map(Ok),
-                capacity,
-                max_length_check,
-                validator,
-                extra,
-                definitions,
-                recursion_guard,
-            ),
-            Self::PyAny(collection) => validate_iter_to_vec(
-                py,
-                collection.iter()?,
-                capacity,
-                max_length_check,
-                validator,
-                extra,
-                definitions,
-                recursion_guard,
-            ),
-            Self::JsonArray(collection) => validate_iter_to_vec(
-                py,
-                collection.iter().map(Ok),
-                capacity,
-                max_length_check,
-                validator,
-                extra,
-                definitions,
-                recursion_guard,
-            ),
+            Self::List(collection) => validate_set!(collection.iter().map(Ok)),
+            Self::Tuple(collection) => validate_set!(collection.iter().map(Ok)),
+            Self::Set(collection) => validate_set!(collection.iter().map(Ok)),
+            Self::FrozenSet(collection) => validate_set!(collection.iter().map(Ok)),
+            Self::PyAny(collection) => validate_set!(collection.iter()?),
+            Self::JsonArray(collection) => validate_set!(collection.iter().map(Ok)),
         }
     }
 

src/validators/frozenset.rs

@@ -6,17 +6,16 @@ use crate::errors::ValResult;
 use crate::input::Input;
 use crate::recursion_guard::RecursionGuard;
 
-use super::list::{get_items_schema, min_length_check};
+use super::list::min_length_check;
 use super::set::set_build;
 use super::{BuildValidator, CombinedValidator, Definitions, DefinitionsBuilder, Extra, Validator};
 
 #[derive(Debug, Clone)]
 pub struct FrozenSetValidator {
     strict: bool,
-    item_validator: Option<Box<CombinedValidator>>,
+    item_validator: Box<CombinedValidator>,
     min_length: Option<usize>,
     max_length: Option<usize>,
-    generator_max_length: Option<usize>,
     name: String,
 }
 
@@ -34,25 +33,20 @@ impl Validator for FrozenSetValidator {
         definitions: &'data Definitions<CombinedValidator>,
         recursion_guard: &'s mut RecursionGuard,
     ) -> ValResult<'data, PyObject> {
-        let seq = input.validate_frozenset(extra.strict.unwrap_or(self.strict))?;
+        let collection = input.validate_frozenset(extra.strict.unwrap_or(self.strict))?;
 
-        let f_set = match self.item_validator {
-            Some(ref v) => PyFrozenSet::new(
-                py,
-                &seq.validate_to_vec(
-                    py,
-                    input,
-                    self.max_length,
-                    "Frozenset",
-                    v,
-                    extra,
-                    definitions,
-                    recursion_guard,
-                )?,
-            )?,
-            None => PyFrozenSet::new(py, &seq.to_vec(py, input, "Frozenset", self.generator_max_length)?)?,
-        };
-        min_length_check!(input, "Frozenset", self.min_length, f_set);
+        let set = collection.validate_to_set(
+            py,
+            input,
+            self.max_length,
+            "Frozenset",
+            &self.item_validator,
+            extra,
+            definitions,
+            recursion_guard,
+        )?;
+        min_length_check!(input, "Frozenset", self.min_length, set);
+        let f_set = PyFrozenSet::new(py, set)?;
         Ok(f_set.into_py(py))
     }
 
@@ -62,10 +56,7 @@ impl Validator for FrozenSetValidator {
         ultra_strict: bool,
     ) -> bool {
         if ultra_strict {
-            match self.item_validator {
-                Some(ref v) => v.different_strict_behavior(definitions, true),
-                None => false,
-            }
+            self.item_validator.different_strict_behavior(definitions, true)
         } else {
             true
         }
@@ -76,9 +67,6 @@ impl Validator for FrozenSetValidator {
     }
 
     fn complete(&mut self, definitions: &DefinitionsBuilder<CombinedValidator>) -> PyResult<()> {
-        match self.item_validator {
-            Some(ref mut v) => v.complete(definitions),
-            None => Ok(()),
-        }
+        self.item_validator.complete(definitions)
     }
 }