Skip reusing after validators in prebuilt schemas (#1663)

Author: sydney-runkle

src/validators/prebuilt.rs

@@ -17,7 +17,10 @@ impl PrebuiltValidator {
     pub fn try_get_from_schema(type_: &str, schema: &Bound<'_, PyDict>) -> PyResult<Option<CombinedValidator>> {
         get_prebuilt(type_, schema, "__pydantic_validator__", |py_any| {
             let schema_validator = py_any.extract::<Py<SchemaValidator>>()?;
-            if matches!(schema_validator.get().validator, CombinedValidator::FunctionWrap(_)) {
+            if matches!(
+                schema_validator.get().validator,
+                CombinedValidator::FunctionWrap(_) | CombinedValidator::FunctionAfter(_)
+            ) {
                 return Ok(None);
             }
             Ok(Some(Self { schema_validator }.into()))

tests/test_prebuilt.py

@@ -14,11 +14,11 @@ class InnerModel:
         ),
     )
 
-    inner_schema_validator = SchemaValidator(inner_schema)
-    inner_schema_serializer = SchemaSerializer(inner_schema)
+    inner_validator = SchemaValidator(inner_schema)
+    inner_serializer = SchemaSerializer(inner_schema)
     InnerModel.__pydantic_complete__ = True  # pyright: ignore[reportAttributeAccessIssue]
-    InnerModel.__pydantic_validator__ = inner_schema_validator  # pyright: ignore[reportAttributeAccessIssue]
-    InnerModel.__pydantic_serializer__ = inner_schema_serializer  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_validator__ = inner_validator  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_serializer__ = inner_serializer  # pyright: ignore[reportAttributeAccessIssue]
 
     class OuterModel:
         inner: InnerModel
@@ -69,9 +69,9 @@ def serialize_inner(v: InnerModel, serializer) -> Union[dict[str, str], str]:
         serialization=core_schema.wrap_serializer_function_ser_schema(serialize_inner),
     )
 
-    inner_schema_serializer = SchemaSerializer(inner_schema)
+    inner_serializer = SchemaSerializer(inner_schema)
     InnerModel.__pydantic_complete__ = True  # pyright: ignore[reportAttributeAccessIssue]
-    InnerModel.__pydantic_serializer__ = inner_schema_serializer  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_serializer__ = inner_serializer  # pyright: ignore[reportAttributeAccessIssue]
 
     class OuterModel:
         inner: InnerModel
@@ -97,7 +97,6 @@ def __init__(self, inner: InnerModel) -> None:
         ),
     )
 
-    inner_serializer = SchemaSerializer(inner_schema)
     outer_serializer = SchemaSerializer(outer_schema)
 
     # the custom serialization function does apply for the inner model
@@ -130,9 +129,9 @@ def validate_inner(data, validator) -> InnerModel:
         ),
     )
 
-    inner_schema_validator = SchemaValidator(inner_schema)
+    inner_validator = SchemaValidator(inner_schema)
     InnerModel.__pydantic_complete__ = True  # pyright: ignore[reportAttributeAccessIssue]
-    InnerModel.__pydantic_validator__ = inner_schema_validator  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_validator__ = inner_validator  # pyright: ignore[reportAttributeAccessIssue]
 
     class OuterModel:
         inner: InnerModel
@@ -158,7 +157,6 @@ def __init__(self, inner: InnerModel) -> None:
         ),
     )
 
-    inner_validator = SchemaValidator(inner_schema)
     outer_validator = SchemaValidator(outer_schema)
 
     # the custom validation function does apply for the inner model
@@ -170,6 +168,66 @@ def __init__(self, inner: InnerModel) -> None:
     assert result_outer.inner.x == 'hello'
 
 
+def test_prebuilt_not_used_for_after_validator_functions() -> None:
+    class InnerModel:
+        x: str
+
+        def __init__(self, x: str) -> None:
+            self.x = x
+
+    def validate_after(self) -> InnerModel:
+        self.x = self.x + ' modified'
+        return self
+
+    inner_schema = core_schema.no_info_after_validator_function(
+        validate_after,
+        core_schema.model_schema(
+            InnerModel,
+            schema=core_schema.model_fields_schema(
+                {'x': core_schema.model_field(schema=core_schema.str_schema())},
+            ),
+        ),
+    )
+
+    inner_validator = SchemaValidator(inner_schema)
+    InnerModel.__pydantic_complete__ = True  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_validator__ = inner_validator  # pyright: ignore[reportAttributeAccessIssue]
+
+    class OuterModel:
+        inner: InnerModel
+
+        def __init__(self, inner: InnerModel) -> None:
+            self.inner = inner
+
+    outer_schema = core_schema.model_schema(
+        OuterModel,
+        schema=core_schema.model_fields_schema(
+            {
+                'inner': core_schema.model_field(
+                    schema=core_schema.model_schema(
+                        InnerModel,
+                        schema=core_schema.model_fields_schema(
+                            # note, we use a simple str schema (with no custom validation)
+                            # in order to verify that the prebuilt validator from InnerModel is not used
+                            {'x': core_schema.model_field(schema=core_schema.str_schema())},
+                        ),
+                    )
+                )
+            }
+        ),
+    )
+
+    outer_validator = SchemaValidator(outer_schema)
+
+    # the custom validation function does apply for the inner model
+    result_inner = inner_validator.validate_python({'x': 'hello'})
+    assert result_inner.x == 'hello modified'
+
+    # but the outer model doesn't reuse the custom after validator function, so we see simple str val
+    result_outer = outer_validator.validate_python({'inner': {'x': 'hello'}})
+    assert result_outer.inner.x == 'hello'
+
+
 def test_reuse_plain_serializer_ok() -> None:
     class InnerModel:
         x: str
@@ -188,9 +246,9 @@ def serialize_inner(v: InnerModel) -> str:
         serialization=core_schema.plain_serializer_function_ser_schema(serialize_inner),
     )
 
-    inner_schema_serializer = SchemaSerializer(inner_schema)
+    inner_serializer = SchemaSerializer(inner_schema)
     InnerModel.__pydantic_complete__ = True  # pyright: ignore[reportAttributeAccessIssue]
-    InnerModel.__pydantic_serializer__ = inner_schema_serializer  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_serializer__ = inner_serializer  # pyright: ignore[reportAttributeAccessIssue]
 
     class OuterModel:
         inner: InnerModel
@@ -216,7 +274,6 @@ def __init__(self, inner: InnerModel) -> None:
         ),
     )
 
-    inner_serializer = SchemaSerializer(inner_schema)
     outer_serializer = SchemaSerializer(outer_schema)
 
     # the custom serialization function does apply for the inner model
@@ -243,9 +300,9 @@ def validate_inner(data) -> InnerModel:
 
     inner_schema = core_schema.no_info_plain_validator_function(validate_inner)
 
-    inner_schema_validator = SchemaValidator(inner_schema)
+    inner_validator = SchemaValidator(inner_schema)
     InnerModel.__pydantic_complete__ = True  # pyright: ignore[reportAttributeAccessIssue]
-    InnerModel.__pydantic_validator__ = inner_schema_validator  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_validator__ = inner_validator  # pyright: ignore[reportAttributeAccessIssue]
 
     class OuterModel:
         inner: InnerModel
@@ -271,7 +328,6 @@ def __init__(self, inner: InnerModel) -> None:
         ),
     )
 
-    inner_validator = SchemaValidator(inner_schema)
     outer_validator = SchemaValidator(outer_schema)
 
     # the custom validation function does apply for the inner model
@@ -283,3 +339,67 @@ def __init__(self, inner: InnerModel) -> None:
     result_outer = outer_validator.validate_python({'inner': {'x': 'hello'}})
     assert result_outer.inner.x == 'hello modified'
     assert 'PrebuiltValidator' in repr(outer_validator)
+
+
+def test_reuse_before_validator_ok() -> None:
+    class InnerModel:
+        x: str
+
+        def __init__(self, x: str) -> None:
+            self.x = x
+
+    def validate_before(data) -> dict:
+        data['x'] = data['x'] + ' modified'
+        return data
+
+    inner_schema = core_schema.no_info_before_validator_function(
+        validate_before,
+        core_schema.model_schema(
+            InnerModel,
+            schema=core_schema.model_fields_schema(
+                {'x': core_schema.model_field(schema=core_schema.str_schema())},
+            ),
+        ),
+    )
+
+    inner_validator = SchemaValidator(inner_schema)
+    InnerModel.__pydantic_complete__ = True  # pyright: ignore[reportAttributeAccessIssue]
+    InnerModel.__pydantic_validator__ = inner_validator  # pyright: ignore[reportAttributeAccessIssue]
+
+    class OuterModel:
+        inner: InnerModel
+
+        def __init__(self, inner: InnerModel) -> None:
+            self.inner = inner
+
+    outer_schema = core_schema.model_schema(
+        OuterModel,
+        schema=core_schema.model_fields_schema(
+            {
+                'inner': core_schema.model_field(
+                    schema=core_schema.model_schema(
+                        InnerModel,
+                        schema=core_schema.model_fields_schema(
+                            # note, we use a simple str schema (with no custom validation)
+                            # in order to verify that the prebuilt validator from InnerModel is used instead
+                            {'x': core_schema.model_field(schema=core_schema.str_schema())},
+                        ),
+                    )
+                )
+            }
+        ),
+    )
+
+    outer_validator = SchemaValidator(outer_schema)
+    print(inner_validator)
+    print(outer_validator)
+
+    # the custom validation function does apply for the inner model
+    result_inner = inner_validator.validate_python({'x': 'hello'})
+    assert result_inner.x == 'hello modified'
+    assert 'FunctionBeforeValidator' in repr(inner_validator)
+
+    # the custom validation function does apply for the outer model as well, a before validator is permitted as a prebuilt candidate
+    result_outer = outer_validator.validate_python({'inner': {'x': 'hello'}})
+    assert result_outer.inner.x == 'hello modified'
+    assert 'PrebuiltValidator' in repr(outer_validator)