From ae4cb2851d88955fdcbac7acd3b7d99fe02e9628 Mon Sep 17 00:00:00 2001
From: David Montague <35119617+dmontagu@users.noreply.github.com>
Date: Mon, 24 Apr 2023 10:44:07 -0600
Subject: [PATCH] Fix serialization of models with computed fields (#550)

* Fix serialization of models with computed fields

* Move the n_computed_fields method
---
 src/serializers/computed_fields.rs            |  4 +++
 .../type_serializers/typed_dict.rs            |  9 ++++-
 tests/test.rs                                 | 36 +++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/src/serializers/computed_fields.rs b/src/serializers/computed_fields.rs
index 7e19d1e2f..9da9f5c3c 100644
--- a/src/serializers/computed_fields.rs
+++ b/src/serializers/computed_fields.rs
@@ -32,6 +32,10 @@ impl ComputedFields {
         }
     }
 
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
     pub fn to_python(
         &self,
         model: &PyAny,
diff --git a/src/serializers/type_serializers/typed_dict.rs b/src/serializers/type_serializers/typed_dict.rs
index 445ac7845..1f799030f 100644
--- a/src/serializers/type_serializers/typed_dict.rs
+++ b/src/serializers/type_serializers/typed_dict.rs
@@ -144,6 +144,13 @@ impl TypedDictSerializer {
         }
     }
 
+    pub fn n_computed_fields(&self) -> usize {
+        match self.computed_fields {
+            None => 0,
+            Some(ref computed_fields) => computed_fields.len(),
+        }
+    }
+
     fn exclude_default(&self, value: &PyAny, extra: &Extra, field: &TypedDictField) -> PyResult<bool> {
         if extra.exclude_defaults {
             if let Some(default) = field.serializer.get_default(value.py())? {
@@ -262,7 +269,7 @@ impl TypeSerializer for TypedDictSerializer {
                 };
                 let expected_len = match self.include_extra {
                     true => py_dict.len(),
-                    false => self.fields.len(),
+                    false => self.fields.len() + self.n_computed_fields(),
                 };
                 // NOTE! As above, we maintain the order of the input dict assuming that's right
                 // we don't both with `used_fields` here because on unions, `to_python(..., mode='json')` is used
diff --git a/tests/test.rs b/tests/test.rs
index dc2fd5a5c..95e9f4a62 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -28,4 +28,40 @@ mod tests {
             SchemaSerializer::py_new(py, schema, None).unwrap();
         })
     }
+
+    #[test]
+    fn test_serialize_computed_fields() {
+        Python::with_gil(|py| {
+            let code = r#"
+class A:
+    @property
+    def b(self) -> str:
+        return "b"
+
+schema = {
+    "cls": A,
+    "config": {},
+    "schema": {
+        "computed_fields": [{"property_name": "b", "type": "computed-field"}],
+        "fields": {},
+        "return_fields_set": True,
+        "type": "typed-dict",
+    },
+    "type": "model",
+}
+a = A()
+            "#;
+            let locals = PyDict::new(py);
+            py.run(code, None, Some(locals)).unwrap();
+            let a: &PyAny = locals.get_item("a").unwrap().extract().unwrap();
+            let schema: &PyDict = locals.get_item("schema").unwrap().extract().unwrap();
+            let serialized: Vec<u8> = SchemaSerializer::py_new(py, schema, None)
+                .unwrap()
+                .to_json(py, a, None, None, None, true, false, false, false, false, true, None)
+                .unwrap()
+                .extract(py)
+                .unwrap();
+            assert_eq!(serialized, b"{\"b\":\"b\"}");
+        })
+    }
 }