Support for typing.NotRequired (#17)

Author: giograno

src/py_avro_schema/_schemas.py

@@ -37,6 +37,7 @@
     ForwardRef,
     List,
     Literal,
+    NotRequired,
     Optional,
     Tuple,
     Type,
@@ -1289,11 +1290,15 @@ def _record_field(self, py_field: tuple[str, Type]) -> RecordField:
         aliases, actual_type = get_field_aliases_and_actual_type(py_field[1])
 
         if Option.MARK_NON_TOTAL_TYPED_DICTS in self.options and not self.is_total:
-            # If a TypedDict is marked as total=None, it does not need to contain all the field. However, we need to
+            # If a TypedDict is marked as total=False, it does not need to contain all the field. However, we need to
             # be able to distinguish between the fields that are missing from the ones that are present but set to None.
             # To do that, we extend the original type with str. We will later add a special string
             # (e.g., __td_missing__) as a marker at deserialization time.
             actual_type = Union[actual_type, str]  # type: ignore
+        elif _is_not_required(actual_type):
+            # A field can be marked with typing.NotRequired even in a TypedDict with is not marked with total=False.
+            # Similarly as above, we extend the wrapped type with string.
+            actual_type = Union[_unwrap_not_required(actual_type), str]  # type: ignore
 
         field_obj = RecordField(
             py_type=actual_type,
@@ -1316,6 +1321,16 @@ def _doc_for_class(py_type: Type) -> str:
         return ""
 
 
+def _is_not_required(py_type: Type) -> bool:
+    """Checks if a type is marked with typing.NotRequired"""
+    return get_origin(py_type) is NotRequired  # noqa
+
+
+def _unwrap_not_required(py_type: Type) -> type:
+    """Returns the wrapped type for typing.NotRequired"""
+    return get_args(py_type)[0]
+
+
 def _is_dict_str_any(py_type: Type) -> bool:
     """Return whether a given type is ``Dict[str, Any]``"""
     origin = get_origin(py_type)

tests/test_typed_dict.py

@@ -1,5 +1,5 @@
 from enum import StrEnum
-from typing import Annotated, TypedDict
+from typing import Annotated, NotRequired, TypedDict
 
 import py_avro_schema as pas
 from py_avro_schema._alias import Alias, register_type_alias
@@ -103,13 +103,31 @@ class PyType(TypedDict, total=False):
         "type": "record",
         "name": "PyType",
         "fields": [
-            {
-                "name": "name",
-                "type": "string",
-            },
+            {"name": "name", "type": "string"},
             {"name": "nickname", "type": ["string", "null"]},
             {"name": "age", "type": ["long", "null", "string"]},
             {"name": "opt", "type": [{"namedString": "Opt", "type": "string"}, "null"]},
         ],
     }
     assert_schema(PyType, expected, options=pas.Option.MARK_NON_TOTAL_TYPED_DICTS)
+
+
+def test_non_required_keyword():
+    class PyType(TypedDict):
+        name: str
+        value: NotRequired[str]
+        value_int: NotRequired[int]
+        nullable_value: NotRequired[str | None]
+
+    expected = {
+        "type": "record",
+        "name": "PyType",
+        "fields": [
+            {"name": "name", "type": "string"},
+            {"name": "value", "type": "string"},
+            {"name": "value_int", "type": ["long", "string"]},
+            {"name": "nullable_value", "type": ["string", "null"]},
+        ],
+    }
+
+    assert_schema(PyType, expected, options=pas.Option.MARK_NON_TOTAL_TYPED_DICTS)