feat: allow fields to be frozen and add strict + context support to validate_assignment (#221)

* feat: allow fields to be `frozen`

* add test for `validate_assignment` with `strict` field

* add `strict` and `context` support to `validate_assignment`

Author: PrettyWood

pydantic_core/_pydantic_core.pyi

@@ -21,7 +21,9 @@ class SchemaValidator:
     def isinstance_json(
         self, input: 'str | bytes | bytearray', strict: 'bool | None' = None, context: Any = None
     ) -> bool: ...
-    def validate_assignment(self, field: str, input: Any, data: 'dict[str, Any]') -> 'dict[str, Any]': ...
+    def validate_assignment(
+        self, field: str, input: Any, data: 'dict[str, Any]', strict: 'bool | None' = None, context: Any = None
+    ) -> 'dict[str, Any]': ...
 
 class SchemaError(Exception):
     pass

pydantic_core/_types.py

@@ -125,6 +125,7 @@ class TypedDictField(TypedDict, total=False):
     default_factory: Callable[[], Any]
     on_error: Literal['raise', 'omit', 'fallback_on_default']  # default: 'raise'
     alias: Union[str, List[Union[str, int]], List[List[Union[str, int]]]]
+    frozen: bool
 
 
 class TypedDictSchema(TypedDict, total=False):

src/errors/kinds.rs

@@ -30,6 +30,8 @@ pub enum ErrorKind {
     DictAttributesType,
     #[strum(message = "Field required")]
     Missing,
+    #[strum(message = "Field is frozen")]
+    Frozen,
     #[strum(message = "Extra inputs are not permitted")]
     ExtraForbidden,
     #[strum(message = "Keys should be strings")]

src/validators/mod.rs

@@ -175,12 +175,20 @@ impl SchemaValidator {
         }
     }
 
-    pub fn validate_assignment(&self, py: Python, field: String, input: &PyAny, data: &PyDict) -> PyResult<PyObject> {
+    pub fn validate_assignment(
+        &self,
+        py: Python,
+        field: String,
+        input: &PyAny,
+        data: &PyDict,
+        strict: Option<bool>,
+        context: Option<&PyAny>,
+    ) -> PyResult<PyObject> {
         let extra = Extra {
             data: Some(data),
             field: Some(field.as_str()),
-            strict: None,
-            context: None,
+            strict,
+            context,
         };
         let r = self
             .validator

src/validators/typed_dict.rs

@@ -32,6 +32,7 @@ struct TypedDictField {
     default: Option<PyObject>,
     default_factory: Option<PyObject>,
     validator: CombinedValidator,
+    frozen: bool,
 }
 
 impl TypedDictField {
@@ -182,8 +183,10 @@ impl BuildValidator for TypedDictValidator {
                 default,
                 default_factory,
                 on_error,
+                frozen: field_info.get_as::<bool>(intern!(py, "frozen"))?.unwrap_or(false),
             });
         }
+
         Ok(Self {
             fields,
             check_extra,
@@ -433,7 +436,11 @@ impl TypedDictValidator {
         };
 
         if let Some(field) = self.fields.iter().find(|f| f.name == field) {
-            prepare_result(field.validator.validate(py, input, extra, slots, recursion_guard))
+            if field.frozen {
+                Err(ValError::new_with_loc(ErrorKind::Frozen, input, field.name.to_string()))
+            } else {
+                prepare_result(field.validator.validate(py, input, extra, slots, recursion_guard))
+            }
         } else if self.check_extra && !self.forbid_extra {
             // this is the "allow" case of extra_behavior
             match self.extra_validator {

tests/test_validation_context.py

@@ -1,6 +1,6 @@
 import pytest
 
-from pydantic_core import ValidationError
+from pydantic_core import SchemaValidator, ValidationError
 
 from .conftest import PyAndJson
 
@@ -89,3 +89,29 @@ def f(input_value, *, validator, context, **kwargs):
         v.isinstance_test('foobar')
 
     assert v.isinstance_test('foobar', None, {'error'}) is False
+
+
+def test_validate_assignment_with_context():
+    def f1(input_value, *, context, **kwargs):
+        context['f1'] = input_value
+        return input_value + f'| context: {context}'
+
+    def f2(input_value, *, context, **kwargs):
+        context['f2'] = input_value
+        return input_value + f'| context: {context}'
+
+    v = SchemaValidator(
+        {
+            'type': 'typed-dict',
+            'fields': {
+                'f1': {'schema': {'type': 'function', 'mode': 'plain', 'function': f1}},
+                'f2': {'schema': {'type': 'function', 'mode': 'plain', 'function': f2}},
+            },
+        }
+    )
+
+    m1 = v.validate_python({'f1': '1', 'f2': '2'}, None, {'x': 'y'})
+    assert m1 == {'f1': "1| context: {'x': 'y', 'f1': '1'}", 'f2': "2| context: {'x': 'y', 'f1': '1', 'f2': '2'}"}
+
+    m2 = v.validate_assignment('f1', '3', m1, None, {'x': 'y'})
+    assert m2 == {'f1': "3| context: {'x': 'y', 'f1': '3'}", 'f2': "2| context: {'x': 'y', 'f1': '1', 'f2': '2'}"}

tests/validators/test_typed_dict.py

@@ -241,6 +241,24 @@ def test_validate_assignment():
     assert v.validate_assignment('field_a', b'abc', {'field_a': 'test'}) == ({'field_a': 'abc'}, {'field_a'})
 
 
+def test_validate_assignment_strict_field():
+    v = SchemaValidator(
+        {
+            'type': 'typed-dict',
+            'return_fields_set': True,
+            'fields': {'field_a': {'schema': {'type': 'str', 'strict': True}}},
+        }
+    )
+
+    assert v.validate_python({'field_a': 'test'}) == ({'field_a': 'test'}, {'field_a'})
+
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_assignment('field_a', b'abc', {'field_a': 'test'})
+    assert exc_info.value.errors() == [
+        {'input_value': b'abc', 'kind': 'str_type', 'loc': ['field_a'], 'message': 'Input should be a valid string'}
+    ]
+
+
 def test_validate_assignment_functions():
     calls = []
 
@@ -336,6 +354,24 @@ def test_validate_assignment_allow_extra_validate():
     ]
 
 
+def test_validate_assignment_with_strict():
+    v = SchemaValidator(
+        {'type': 'typed-dict', 'fields': {'x': {'schema': {'type': 'str'}}, 'y': {'schema': {'type': 'int'}}}}
+    )
+
+    r = v.validate_python({'x': 'a', 'y': '123'})
+    assert r == {'x': 'a', 'y': 123}
+
+    assert v.validate_assignment('y', '124', r) == {'x': 'a', 'y': 124}
+
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_assignment('y', '124', r, True)
+
+    assert exc_info.value.errors() == [
+        {'kind': 'int_type', 'loc': ['y'], 'message': 'Input should be a valid integer', 'input_value': '124'}
+    ]
+
+
 def test_json_error():
     v = SchemaValidator(
         {'type': 'typed-dict', 'fields': {'field_a': {'schema': {'type': 'list', 'items_schema': 'int'}}}}
@@ -1211,3 +1247,25 @@ def wrap_function(input_value, *, validator, **kwargs):
         assert v.validate_test({'x': ['foo']}) == {'x': '1'}
         assert v.validate_test({'x': ['foo', 'bar']}) == {'x': '2'}
         assert v.validate_test({'x': {'a': 'b'}}) == {'x': "{'a': 'b'}"}
+
+
+def test_frozen_field():
+    v = SchemaValidator(
+        {
+            'type': 'typed-dict',
+            'fields': {
+                'name': {'schema': {'type': 'str'}},
+                'age': {'schema': {'type': 'int'}},
+                'is_developer': {'schema': {'type': 'bool'}, 'default': True, 'frozen': True},
+            },
+        }
+    )
+    r1 = v.validate_python({'name': 'Samuel', 'age': '36'})
+    assert r1 == {'name': 'Samuel', 'age': 36, 'is_developer': True}
+    r2 = v.validate_assignment('age', '35', r1)
+    assert r2 == {'name': 'Samuel', 'age': 35, 'is_developer': True}
+    with pytest.raises(ValidationError) as exc_info:
+        v.validate_assignment('is_developer', False, r2)
+    assert exc_info.value.errors() == [
+        {'kind': 'frozen', 'loc': ['is_developer'], 'message': 'Field is frozen', 'input_value': False}
+    ]