add Forward Ref optional typing

Author: bentsku

src/py_avro_schema/_testing.py

@@ -16,15 +16,15 @@
 
 import dataclasses
 import difflib
-from typing import Dict, Type, Union
+from typing import Dict, List, Type, Union
 
 import avro.schema  # type: ignore
 import orjson
 
 import py_avro_schema._schemas
 
 
-def assert_schema(py_type: Type, expected_schema: Union[str, Dict[str, str]], **kwargs) -> None:
+def assert_schema(py_type: Type, expected_schema: Union[str, Dict[str, str], List[str]], **kwargs) -> None:
     """Test that the given Python type results in the correct Avro schema"""
     if not kwargs.pop("do_auto_namespace", False):
         kwargs["options"] = kwargs.get("options", py_avro_schema.Option(0)) | py_avro_schema.Option.NO_AUTO_NAMESPACE

tests/test_primitives.py

@@ -29,6 +29,7 @@
 
 import py_avro_schema as pas
 import py_avro_schema._schemas
+from py_avro_schema._testing import PyType as TestPyType
 from py_avro_schema._testing import assert_schema
 
 
@@ -418,6 +419,67 @@ def test_optional_str_py310():
     assert_schema(py_type, expected)
 
 
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="Requires Python 3.10+")
+def test_optional_forward_ref_py310():
+    class PyType:
+        forward_ref: "TestPyType | None"
+
+    expected = {
+        "fields": [
+            {
+                "name": "forward_ref",
+                "type": [
+                    "PyType",
+                    "null",
+                ],
+            },
+        ],
+        "name": "PyType",
+        "type": "record",
+    }
+    assert_schema(PyType, expected)
+
+
+def test_optional_forward_ref_with_union():
+    class PyType:
+        forward_ref: Union["TestPyType", None]
+
+    expected = {
+        "fields": [
+            {
+                "name": "forward_ref",
+                "type": [
+                    "PyType",
+                    "null",
+                ],
+            },
+        ],
+        "name": "PyType",
+        "type": "record",
+    }
+    assert_schema(PyType, expected)
+
+
+def test_optional_forward_ref():
+    class PyType:
+        forward_ref: Optional["TestPyType"]
+
+    expected = {
+        "fields": [
+            {
+                "name": "forward_ref",
+                "type": [
+                    "PyType",
+                    "null",
+                ],
+            },
+        ],
+        "name": "PyType",
+        "type": "record",
+    }
+    assert_schema(PyType, expected)
+
+
 def test_enum():
     class PyType(enum.Enum):
         RED = "RED"