support __pydantic_extra__ when serializing

Author: samuelcolvin

src/serializers/shared.rs

@@ -3,7 +3,7 @@ use std::fmt::Debug;
 
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PySet};
+use pyo3::types::{PyDict, PyMapping, PySet};
 use pyo3::{intern, PyTraverseError, PyVisit};
 
 use enum_dispatch::enum_dispatch;
@@ -329,17 +329,36 @@ pub(super) fn object_to_dict<'py>(value: &'py PyAny, is_model: bool, extra: &Ext
     let py = value.py();
     let attr = value.getattr(intern!(py, "__dict__"))?;
     let attrs: &PyDict = attr.downcast()?;
-    if is_model && extra.exclude_unset {
-        let fields_set: &PySet = value.getattr(intern!(py, "__pydantic_fields_set__"))?.downcast()?;
+    if is_model {
+        if let Some(extra) = get_pydantic_extra(value) {
+            attrs.update(extra)?;
+        }
+
+        if extra.exclude_unset {
+            let fields_set: &PySet = value.getattr(intern!(py, "__pydantic_fields_set__"))?.downcast()?;
 
-        let new_attrs = attrs.copy()?;
-        for key in new_attrs.keys() {
-            if !fields_set.contains(key)? {
-                new_attrs.del_item(key)?;
+            let new_attrs = attrs.copy()?;
+            for key in new_attrs.keys() {
+                if !fields_set.contains(key)? {
+                    new_attrs.del_item(key)?;
+                }
             }
+            Ok(new_attrs)
+        } else {
+            Ok(attrs)
         }
-        Ok(new_attrs)
     } else {
         Ok(attrs)
     }
 }
+
+fn get_pydantic_extra(value: &PyAny) -> Option<&PyMapping> {
+    let extra = match value.getattr(intern!(value.py(), "__pydantic_extra__")) {
+        Ok(extra) => extra,
+        Err(_) => return None,
+    };
+    match extra.downcast::<PyMapping>() {
+        Ok(attrs) => Some(attrs),
+        Err(_) => None,
+    }
+}

tests/serializers/test_model.py

@@ -741,3 +741,31 @@ def random_n(self) -> int:
     del sq.area
     assert s.to_python(sq, by_alias=False) == {'side': 0, 'area': 0, 'random_n': the_random_n}
     assert s.to_python(sq, exclude={'random_n'}) == {'side': 0, 'area': 0}
+
+
+def test_extra():
+    class MyModel:
+        # this is not required, but it avoids `__pydantic_fields_set__` being included in `__dict__`
+        __slots__ = '__dict__', '__pydantic_extra__', '__pydantic_fields_set__'
+        field_a: str
+        field_b: int
+
+    schema = 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='allow',
+        ),
+    )
+    v = SchemaValidator(schema)
+    m = v.validate_python({'field_a': 'test', 'field_b': 12, 'field_c': 'extra'})
+    assert isinstance(m, MyModel)
+    assert m.__dict__ == {'field_a': 'test', 'field_b': 12}
+    assert m.__pydantic_extra__ == {'field_c': 'extra'}
+    assert m.__pydantic_fields_set__ == {'field_a', 'field_b', 'field_c'}
+
+    s = SchemaSerializer(schema)
+    assert s.to_python(m) == {'field_a': 'test', 'field_b': 12, 'field_c': 'extra'}