Support annotated pydantic model

jpmorganchase · Nov 17, 2023 · 8aa4383 · 8aa4383
1 parent 1747a4c
commit 8aa4383
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/src/py_avro_schema/_schemas.py b/src/py_avro_schema/_schemas.py
@@ -872,6 +872,7 @@ class PydanticSchema(RecordSchema):
     @classmethod
     def handles_type(cls, py_type: Type) -> bool:
         """Whether this schema class can represent a given Python class"""
+        py_type = _type_from_annotated(py_type)
         return hasattr(py_type, "__fields__")
 
     def __init__(self, py_type: Type[pydantic.BaseModel], namespace: Optional[str] = None, options: Option = Option(0)):

diff --git a/tests/test_pydantic.py b/tests/test_pydantic.py
@@ -38,6 +38,23 @@ class PyType(pydantic.BaseModel):
     assert_schema(PyType, expected)
 
 
+def test_class_annotated():
+    class PyType(pydantic.BaseModel):
+        field_a: str
+
+    expected = {
+        "type": "record",
+        "name": "PyType",
+        "fields": [
+            {
+                "name": "field_a",
+                "type": "string",
+            }
+        ],
+    }
+    assert_schema(Annotated[PyType, ...], expected)
+
+
 def test_string_field_default():
     class PyType(pydantic.BaseModel):
         field_a: str = ""