Support revalidation of parametrized generics (#1489)

Author: sydney-runkle

python/pydantic_core/core_schema.py

@@ -3056,6 +3056,7 @@ def model_fields_schema(
 class ModelSchema(TypedDict, total=False):
     type: Required[Literal['model']]
     cls: Required[Type[Any]]
+    generic_origin: Type[Any]
     schema: Required[CoreSchema]
     custom_init: bool
     root_model: bool
@@ -3074,6 +3075,7 @@ def model_schema(
     cls: Type[Any],
     schema: CoreSchema,
     *,
+    generic_origin: Type[Any] | None = None,
     custom_init: bool | None = None,
     root_model: bool | None = None,
     post_init: str | None = None,
@@ -3120,6 +3122,8 @@ class MyModel:
     Args:
         cls: The class to use for the model
         schema: The schema to use for the model
+        generic_origin: The origin type used for this model, if it's a parametrized generic. Ex,
+            if this model schema represents `SomeModel[int]`, generic_origin is `SomeModel`
         custom_init: Whether the model has a custom init method
         root_model: Whether the model is a `RootModel`
         post_init: The call after init to use for the model
@@ -3136,6 +3140,7 @@ class MyModel:
     return _dict_not_none(
         type='model',
         cls=cls,
+        generic_origin=generic_origin,
         schema=schema,
         custom_init=custom_init,
         root_model=root_model,
@@ -3289,6 +3294,7 @@ def dataclass_args_schema(
 class DataclassSchema(TypedDict, total=False):
     type: Required[Literal['dataclass']]
     cls: Required[Type[Any]]
+    generic_origin: Type[Any]
     schema: Required[CoreSchema]
     fields: Required[List[str]]
     cls_name: str
@@ -3308,6 +3314,7 @@ def dataclass_schema(
     schema: CoreSchema,
     fields: List[str],
     *,
+    generic_origin: Type[Any] | None = None,
     cls_name: str | None = None,
     post_init: bool | None = None,
     revalidate_instances: Literal['always', 'never', 'subclass-instances'] | None = None,
@@ -3328,6 +3335,8 @@ def dataclass_schema(
         schema: The schema to use for the dataclass fields
         fields: Fields of the dataclass, this is used in serialization and in validation during re-validation
             and while validating assignment
+        generic_origin: The origin type used for this dataclass, if it's a parametrized generic. Ex,
+            if this model schema represents `SomeDataclass[int]`, generic_origin is `SomeDataclass`
         cls_name: The name to use in error locs, etc; this is useful for generics (default: `cls.__name__`)
         post_init: Whether to call `__post_init__` after validation
         revalidate_instances: whether instances of models and dataclasses (including subclass instances)
@@ -3343,6 +3352,7 @@ def dataclass_schema(
     return _dict_not_none(
         type='dataclass',
         cls=cls,
+        generic_origin=generic_origin,
         fields=fields,
         cls_name=cls_name,
         schema=schema,

src/validators/dataclass.rs

@@ -438,6 +438,7 @@ pub struct DataclassValidator {
     strict: bool,
     validator: Box<CombinedValidator>,
     class: Py<PyType>,
+    generic_origin: Option<Py<PyType>>,
     fields: Vec<Py<PyString>>,
     post_init: Option<Py<PyString>>,
     revalidate: Revalidate,
@@ -461,6 +462,7 @@ impl BuildValidator for DataclassValidator {
         let config = config.as_ref();
 
         let class = schema.get_as_req::<Bound<'_, PyType>>(intern!(py, "cls"))?;
+        let generic_origin: Option<Bound<'_, PyType>> = schema.get_as(intern!(py, "generic_origin"))?;
         let name = match schema.get_as_req::<String>(intern!(py, "cls_name")) {
             Ok(name) => name,
             Err(_) => class.getattr(intern!(py, "__name__"))?.extract()?,
@@ -480,6 +482,7 @@ impl BuildValidator for DataclassValidator {
             strict: is_strict(schema, config)?,
             validator: Box::new(validator),
             class: class.into(),
+            generic_origin: generic_origin.map(std::convert::Into::into),
             fields,
             post_init,
             revalidate: Revalidate::from_str(
@@ -496,7 +499,11 @@ impl BuildValidator for DataclassValidator {
     }
 }
 
-impl_py_gc_traverse!(DataclassValidator { class, validator });
+impl_py_gc_traverse!(DataclassValidator {
+    class,
+    generic_origin,
+    validator
+});
 
 impl Validator for DataclassValidator {
     fn validate<'py>(
@@ -510,10 +517,30 @@ impl Validator for DataclassValidator {
             return self.validate_init(py, self_instance, input, state);
         }
 
-        // same logic as on models
+        // same logic as on models, see more explicit comment in model.rs
         let class = self.class.bind(py);
-        if let Some(py_input) = input_as_python_instance(input, class) {
-            if self.revalidate.should_revalidate(py_input, class) {
+        let generic_origin_class = self.generic_origin.as_ref().map(|go| go.bind(py));
+
+        let (py_instance_input, force_revalidate): (Option<&Bound<'_, PyAny>>, bool) =
+            match input_as_python_instance(input, class) {
+                Some(x) => (Some(x), false),
+                None => {
+                    // if the model has a generic origin, we allow input data to be instances of the generic origin rather than the class,
+                    // as cases like isinstance(SomeModel[Int], SomeModel[Any]) fail the isinstance check, but are valid, we just have to enforce
+                    // that the data is revalidated, hence we set force_revalidate to true
+                    if generic_origin_class.is_some() {
+                        match input_as_python_instance(input, generic_origin_class.unwrap()) {
+                            Some(x) => (Some(x), true),
+                            None => (None, false),
+                        }
+                    } else {
+                        (None, false)
+                    }
+                }
+            };
+
+        if let Some(py_input) = py_instance_input {
+            if self.revalidate.should_revalidate(py_input, class) || force_revalidate {
                 let input_dict = self.dataclass_to_dict(py_input)?;
                 let val_output = self.validator.validate(py, input_dict.as_any(), state)?;
                 let dc = create_class(self.class.bind(py))?;

src/validators/model.rs

@@ -55,6 +55,7 @@ pub struct ModelValidator {
     revalidate: Revalidate,
     validator: Box<CombinedValidator>,
     class: Py<PyType>,
+    generic_origin: Option<Py<PyType>>,
     post_init: Option<Py<PyString>>,
     frozen: bool,
     custom_init: bool,
@@ -76,6 +77,7 @@ impl BuildValidator for ModelValidator {
         let config = schema.get_as(intern!(py, "config"))?;
 
         let class: Bound<'_, PyType> = schema.get_as_req(intern!(py, "cls"))?;
+        let generic_origin: Option<Bound<'_, PyType>> = schema.get_as(intern!(py, "generic_origin"))?;
         let sub_schema = schema.get_as_req(intern!(py, "schema"))?;
         let validator = build_validator(&sub_schema, config.as_ref(), definitions)?;
         let name = class.getattr(intern!(py, "__name__"))?.extract()?;
@@ -93,6 +95,7 @@ impl BuildValidator for ModelValidator {
             )?,
             validator: Box::new(validator),
             class: class.into(),
+            generic_origin: generic_origin.map(std::convert::Into::into),
             post_init: schema.get_as(intern!(py, "post_init"))?,
             frozen: schema.get_as(intern!(py, "frozen"))?.unwrap_or(false),
             custom_init: schema.get_as(intern!(py, "custom_init"))?.unwrap_or(false),
@@ -105,7 +108,11 @@ impl BuildValidator for ModelValidator {
     }
 }
 
-impl_py_gc_traverse!(ModelValidator { class, validator });
+impl_py_gc_traverse!(ModelValidator {
+    class,
+    generic_origin,
+    validator
+});
 
 impl Validator for ModelValidator {
     fn validate<'py>(
@@ -119,13 +126,33 @@ impl Validator for ModelValidator {
             return self.validate_init(py, self_instance, input, state);
         }
 
+        let class = self.class.bind(py);
+        let generic_origin_class = self.generic_origin.as_ref().map(|go| go.bind(py));
+
         // if we're in strict mode, we require an exact instance of the class (from python, with JSON an object is ok)
         // if we're not in strict mode, instances subclasses are okay, as well as dicts, mappings, from attributes etc.
         // if the input is an instance of the class, we "revalidate" it - e.g. we extract and reuse `__pydantic_fields_set__`
         // but use from attributes to create a new instance of the model field type
-        let class = self.class.bind(py);
-        if let Some(py_input) = input_as_python_instance(input, class) {
-            if self.revalidate.should_revalidate(py_input, class) {
+        let (py_instance_input, force_revalidate): (Option<&Bound<'_, PyAny>>, bool) =
+            match input_as_python_instance(input, class) {
+                Some(x) => (Some(x), false),
+                None => {
+                    // if the model has a generic origin, we allow input data to be instances of the generic origin rather than the class,
+                    // as cases like isinstance(SomeModel[Int], SomeModel[Any]) fail the isinstance check, but are valid, we just have to enforce
+                    // that the data is revalidated, hence we set force_revalidate to true
+                    if generic_origin_class.is_some() {
+                        match input_as_python_instance(input, generic_origin_class.unwrap()) {
+                            Some(x) => (Some(x), true),
+                            None => (None, false),
+                        }
+                    } else {
+                        (None, false)
+                    }
+                }
+            };
+
+        if let Some(py_input) = py_instance_input {
+            if self.revalidate.should_revalidate(py_input, class) || force_revalidate {
                 let fields_set = py_input.getattr(intern!(py, DUNDER_FIELDS_SET_KEY))?;
                 if self.root_model {
                     let inner_input = py_input.getattr(intern!(py, ROOT_FIELD))?;