From c150587b875caf3c57797396a5143ccd9697963a Mon Sep 17 00:00:00 2001
From: David Hewitt <mail@davidhewitt.dev>
Date: Fri, 19 Apr 2024 17:01:25 +0100
Subject: [PATCH] fix interaction between `extra != 'ignore'` and
 `from_attributes=True`

---
 src/input/input_python.rs      |  6 ++-
 src/validators/typed_dict.rs   | 10 ++--
 tests/validators/test_model.py | 93 ++++++++++++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 5 deletions(-)

diff --git a/src/input/input_python.rs b/src/input/input_python.rs
index 0c9a807b1..b2284efb5 100644
--- a/src/input/input_python.rs
+++ b/src/input/input_python.rs
@@ -622,7 +622,7 @@ fn from_attributes_applicable(obj: &Bound<'_, PyAny>) -> bool {
     // I don't think it's a very good list at all! But it doesn't have to be at perfect, it just needs to avoid
     // the most egregious foot guns, it's mostly just to catch "builtins"
     // still happy to add more or do something completely different if anyone has a better idea???
-    // dbg!(obj, module_name);
+    // dbg!(obj, &module_name);
     !matches!(module_name.to_str(), Ok("builtins" | "datetime" | "collections"))
 }
 
@@ -808,6 +808,10 @@ impl<'py> ValidatedDict<'py> for GenericPyMapping<'_, 'py> {
         }
     }
 
+    fn is_py_get_attr(&self) -> bool {
+        matches!(self, Self::GetAttr(..))
+    }
+
     fn as_py_dict(&self) -> Option<&Bound<'py, PyDict>> {
         match self {
             Self::Dict(dict) => Some(dict),
diff --git a/src/validators/typed_dict.rs b/src/validators/typed_dict.rs
index c90196c7c..5aa8a5f00 100644
--- a/src/validators/typed_dict.rs
+++ b/src/validators/typed_dict.rs
@@ -156,10 +156,12 @@ impl Validator for TypedDictValidator {
 
         // we only care about which keys have been used if we're iterating over the object for extra after
         // the first pass
-        let mut used_keys: Option<AHashSet<&str>> = match self.extra_behavior {
-            ExtraBehavior::Allow | ExtraBehavior::Forbid => Some(AHashSet::with_capacity(self.fields.len())),
-            ExtraBehavior::Ignore => None,
-        };
+        let mut used_keys: Option<AHashSet<&str>> =
+            if self.extra_behavior == ExtraBehavior::Ignore || dict.is_py_get_attr() {
+                None
+            } else {
+                Some(AHashSet::with_capacity(self.fields.len()))
+            };
 
         {
             let state = &mut state.rebind_extra(|extra| extra.data = Some(output_dict.clone()));
diff --git a/tests/validators/test_model.py b/tests/validators/test_model.py
index 69b6a06ec..bdcfbb2e9 100644
--- a/tests/validators/test_model.py
+++ b/tests/validators/test_model.py
@@ -72,6 +72,99 @@ class MyModel:
     assert m.__dict__ == {'field_a': 'test', 'field_b': 12}
 
 
+def test_model_class_extra_forbid():
+    class MyModel:
+        class Meta:
+            pass
+
+        # this is not required, but it avoids `__pydantic_fields_set__` being included in `__dict__`
+        __slots__ = '__dict__', '__pydantic_fields_set__', '__pydantic_extra__', '__pydantic_private__'
+        field_a: str
+        field_b: int
+
+    class Wrapper:
+        def __init__(self, inner):
+            self._inner = inner
+
+        def __dir__(self):
+            return dir(self._inner)
+
+        def __getattr__(self, key):
+            return getattr(self._inner, key)
+
+    v = SchemaValidator(
+        core_schema.model_schema(
+            MyModel,
+            core_schema.model_fields_schema(
+                {
+                    'field_a': core_schema.model_field(core_schema.str_schema()),
+                    'field_b': core_schema.model_field(core_schema.int_schema()),
+                },
+                extra_behavior='forbid',
+            ),
+        )
+    )
+    m = v.validate_python({'field_a': 'test', 'field_b': 12})
+    assert isinstance(m, MyModel)
+    assert m.field_a == 'test'
+    assert m.field_b == 12
+
+    # try revalidating from the model's attributes
+    m = v.validate_python(Wrapper(m), from_attributes=True)
+
+    with pytest.raises(ValidationError) as exc_info:
+        m = v.validate_python({'field_a': 'test', 'field_b': 12, 'field_c': 'extra'})
+
+    assert exc_info.value.errors(include_url=False) == [
+        {'type': 'extra_forbidden', 'loc': ('field_c',), 'msg': 'Extra inputs are not permitted', 'input': 'extra'}
+    ]
+
+
+@pytest.mark.parametrize('extra_behavior', ['allow', 'ignore', 'forbid'])
+def test_model_class_extra_forbid_from_attributes(extra_behavior: str):
+    # iterating attributes includes much more than just __dict__, so need
+    # careful interaction with __extra__
+
+    class MyModel:
+        # this is not required, but it avoids `__pydantic_fields_set__` being included in `__dict__`
+        __slots__ = '__dict__', '__pydantic_fields_set__', '__pydantic_extra__', '__pydantic_private__'
+        field_a: str
+        field_b: int
+
+    class Data:
+        # https://github.com/pydantic/pydantic/issues/9242
+        class Meta:
+            pass
+
+        def __init__(self, **values):
+            self.__dict__.update(values)
+
+    v = SchemaValidator(
+        core_schema.model_schema(
+            MyModel,
+            core_schema.model_fields_schema(
+                {
+                    'field_a': core_schema.model_field(core_schema.str_schema()),
+                    'field_b': core_schema.model_field(core_schema.int_schema()),
+                },
+                extra_behavior=extra_behavior,
+                from_attributes=True,
+            ),
+        )
+    )
+    m = v.validate_python(Data(field_a='test', field_b=12))
+    assert isinstance(m, MyModel)
+    assert m.field_a == 'test'
+    assert m.field_b == 12
+
+    # with from_attributes, extra is basically ignored
+    m = v.validate_python(Data(field_a='test', field_b=12, field_c='extra'))
+    assert isinstance(m, MyModel)
+    assert m.field_a == 'test'
+    assert m.field_b == 12
+    assert not hasattr(m, 'field_c')
+
+
 def test_model_class_setattr():
     setattr_calls = []