feat(pydantic): allow pydantic model instances as defaults (#64) (#65)

jpmorganchase · Mar 18, 2024 · d7289db · d7289db
2 parents 40687f5 + b841268
commit d7289db
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 9 deletions.
diff --git a/.flake8 b/.flake8
@@ -13,7 +13,7 @@
 [flake8]
 # Flake 8 linting configuration, not supported in pyproject.toml
 
-ignore = E203,W503
+ignore = E203,W503,E701
 max-line-length = 120
 exclude =
   .svn,

diff --git a/src/py_avro_schema/_schemas.py b/src/py_avro_schema/_schemas.py
@@ -619,8 +619,9 @@ def handles_type(cls, py_type: Type) -> bool:
 
         # Support for `X | Y` syntax available in Python 3.10+
         # equivalent to `typing.Union[X, Y]`
-        if getattr(types, "UnionType", None):
-            return origin == Union or origin == types.UnionType  # noqa: E721
+        union_type = getattr(types, "UnionType", None)
+        if union_type:
+            return origin == Union or origin == union_type
         return origin == Union
 
     def __init__(self, py_type: Type[Union[Any]], namespace: Optional[str] = None, options: Option = Option(0)):
@@ -862,6 +863,7 @@ def _record_field(self, py_field: dataclasses.Field) -> RecordField:
             default=default,
             options=self.options,
         )
+
         return field_obj
 
 
@@ -901,6 +903,10 @@ def _record_field(self, name: str, py_field: pydantic.fields.FieldInfo) -> Recor
         )
         return field_obj
 
+    def make_default(self, py_default: pydantic.BaseModel) -> JSONObj:
+        """Return an Avro schema compliant default value for a given Python value"""
+        return {key: _schema_obj(value.__class__).make_default(value) for key, value in py_default}
+
     def _annotation(self, field_name: str) -> Type:
         """
         Fetch the raw annotation for a given field name

diff --git a/tests/test_primitives.py b/tests/test_primitives.py
@@ -44,8 +44,7 @@ def test_str_annotated():
 
 
 def test_str_subclass():
-    class PyType(str):
-        ...
+    class PyType(str): ...
 
     expected = {
         "type": "string",
@@ -55,8 +54,7 @@ class PyType(str):
 
 
 def test_str_subclass_namespaced():
-    class PyType(str):
-        ...
+    class PyType(str): ...
 
     expected = {
         "type": "string",
@@ -68,8 +66,7 @@ class PyType(str):
 def test_str_subclass_other_classes():
     import packaging.version
 
-    class PyType(packaging.version.Version, str):
-        ...
+    class PyType(packaging.version.Version, str): ...
 
     expected = {
         "type": "string",

diff --git a/tests/test_pydantic.py b/tests/test_pydantic.py
@@ -523,3 +523,41 @@ class PyType(Base):
         ],
     }
     assert_schema(PyType, expected)
+
+
+def test_base_model_defaults():
+    class Default(pydantic.BaseModel):
+        field_a: str = "default_a"
+
+    class PyType(pydantic.BaseModel):
+        default: Default = pydantic.Field(..., default_factory=Default)
+
+    class Nested(pydantic.BaseModel):
+        py_type: PyType = pydantic.Field(..., default_factory=PyType)
+
+    expected = {
+        "fields": [
+            {
+                "default": {"default": {"field_a": "default_a"}},
+                "name": "py_type",
+                "type": {
+                    "fields": [
+                        {
+                            "default": {"field_a": "default_a"},
+                            "name": "default",
+                            "type": {
+                                "fields": [{"default": "default_a", "name": "field_a", "type": "string"}],
+                                "name": "Default",
+                                "type": "record",
+                            },
+                        }
+                    ],
+                    "name": "PyType",
+                    "type": "record",
+                },
+            }
+        ],
+        "name": "Nested",
+        "type": "record",
+    }
+    assert_schema(Nested, expected)