Add support for generic TypedDict following python/cpython#89026

Fatal1ty · Jan 22, 2023 · 985e91a · 985e91a
1 parent 18735dd
commit 985e91a
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 4 deletions.
diff --git a/mashumaro/core/meta/types/pack.py b/mashumaro/core/meta/types/pack.py
@@ -528,7 +528,13 @@ def pack_named_tuple(spec: ValueSpec) -> Expression:
 
 
 def pack_typed_dict(spec: ValueSpec) -> Expression:
-    annotations = spec.type.__annotations__
+    resolved = resolve_type_params(spec.origin_type, get_args(spec.type))[
+        spec.origin_type
+    ]
+    annotations = {
+        k: resolved.get(v, v)
+        for k, v in spec.origin_type.__annotations__.items()
+    }
     all_keys = list(annotations.keys())
     required_keys = getattr(spec.type, "__required_keys__", all_keys)
     optional_keys = getattr(spec.type, "__optional_keys__", [])
@@ -637,7 +643,7 @@ def inner_expr(
             f'{{{inner_expr(0, "key")}: {inner_expr(1, v_type=int)} '
             f"for key, value in {spec.expression}.items()}}"
         )
-    elif is_typed_dict(spec.type):
+    elif is_typed_dict(spec.origin_type):
         return pack_typed_dict(spec)
     elif ensure_generic_mapping(spec, args, typing.Mapping):
         return (

diff --git a/mashumaro/core/meta/types/unpack.py b/mashumaro/core/meta/types/unpack.py
@@ -653,7 +653,13 @@ def unpack_named_tuple(spec: ValueSpec) -> Expression:
 
 
 def unpack_typed_dict(spec: ValueSpec) -> Expression:
-    annotations = spec.type.__annotations__
+    resolved = resolve_type_params(spec.origin_type, get_args(spec.type))[
+        spec.origin_type
+    ]
+    annotations = {
+        k: resolved.get(v, v)
+        for k, v in spec.origin_type.__annotations__.items()
+    }
     all_keys = list(annotations.keys())
     required_keys = getattr(spec.type, "__required_keys__", all_keys)
     optional_keys = getattr(spec.type, "__optional_keys__", [])
@@ -786,7 +792,7 @@ def inner_expr(
             f"{inner_expr(1, v_type=int)} "
             f"for key, value in {spec.expression}.items()}})"
         )
-    elif is_typed_dict(spec.type):
+    elif is_typed_dict(spec.origin_type):
         return unpack_typed_dict(spec)
     elif ensure_generic_mapping(spec, args, typing.Mapping):
         return (

diff --git a/tests/entities.py b/tests/entities.py
@@ -257,6 +257,11 @@ class TypedDictOptionalKeysWithOptional(TypedDict, total=False):
     y: float
 
 
+class GenericTypedDict(TypedDict, Generic[T]):
+    x: T
+    y: int
+
+
 class MyNamedTuple(NamedTuple):
     i: int
     f: float

diff --git a/tests/test_data_types.py b/tests/test_data_types.py
@@ -64,6 +64,7 @@
     DataClassWithoutMixin,
     GenericSerializableList,
     GenericSerializableTypeDataClass,
+    GenericTypedDict,
     MutableString,
     MyDataClass,
     MyDataClassWithUnion,
@@ -1346,6 +1347,26 @@ class DataClass(DataClassDictMixin):
     assert obj.to_dict()
 
 
+def test_unbound_generic_typed_dict():
+    @dataclass
+    class DataClass(DataClassDictMixin):
+        x: GenericTypedDict
+
+    obj = DataClass({"x": "2023-01-22", "y": 42})
+    assert DataClass.from_dict({"x": {"x": "2023-01-22", "y": "42"}}) == obj
+    assert obj.to_dict() == {"x": {"x": "2023-01-22", "y": 42}}
+
+
+def test_bound_generic_typed_dict():
+    @dataclass
+    class DataClass(DataClassDictMixin):
+        x: GenericTypedDict[date]
+
+    obj = DataClass({"x": date(2023, 1, 22), "y": 42})
+    assert DataClass.from_dict({"x": {"x": "2023-01-22", "y": "42"}}) == obj
+    assert obj.to_dict() == {"x": {"x": "2023-01-22", "y": 42}}
+
+
 def test_dataclass_with_init_false_field():
     @dataclass
     class DataClass(DataClassDictMixin):