support generic dataclasses

Author: onursatici

marshmallow_dataclass/__init__.py

@@ -374,7 +374,9 @@ def class_schema(
     >>> class_schema(Custom)().load({})
     Custom(name=None)
     """
-    if not dataclasses.is_dataclass(clazz):
+    if not dataclasses.is_dataclass(clazz) and not _is_generic_alias_of_dataclass(
+        clazz
+    ):
         clazz = dataclasses.dataclass(clazz)
     if not clazz_frame:
         current_frame = inspect.currentframe()
@@ -397,8 +399,7 @@ def _internal_class_schema(
 ) -> Type[marshmallow.Schema]:
     _RECURSION_GUARD.seen_classes[clazz] = clazz.__name__
     try:
-        # noinspection PyDataclass
-        fields: Tuple[dataclasses.Field, ...] = dataclasses.fields(clazz)
+        class_name, fields = _dataclass_name_and_fields(clazz)
     except TypeError:  # Not a dataclass
         try:
             warnings.warn(
@@ -448,7 +449,7 @@ def _internal_class_schema(
         if field.init or include_non_init
     )
 
-    schema_class = type(clazz.__name__, (_base_schema(clazz, base_schema),), attributes)
+    schema_class = type(class_name, (_base_schema(clazz, base_schema),), attributes)
     return cast(Type[marshmallow.Schema], schema_class)
 
 
@@ -812,6 +813,47 @@ def _get_field_default(field: dataclasses.Field):
     return field.default
 
 
+def _is_generic_alias_of_dataclass(clazz: type) -> bool:
+    """
+    Check if given class is a generic alias of a dataclass, if the dataclass is
+    defined as `class A(Generic[T])`, this method will return true if `A[int]` is passed
+    """
+    return typing_inspect.is_generic_type(clazz) and dataclasses.is_dataclass(
+        typing_inspect.get_origin(clazz)
+    )
+
+
+# noinspection PyDataclass
+def _dataclass_name_and_fields(
+    clazz: type,
+) -> Tuple[str, Tuple[dataclasses.Field, ...]]:
+    if not _is_generic_alias_of_dataclass(clazz):
+        return clazz.__name__, dataclasses.fields(clazz)
+
+    base_dataclass = typing_inspect.get_origin(clazz)
+    base_parameters = typing_inspect.get_parameters(base_dataclass)
+    type_arguments = typing_inspect.get_args(clazz)
+    params_to_args = dict(zip(base_parameters, type_arguments))
+    non_generic_fields = [  # swap generic typed fields with types in given type arguments
+        (
+            f.name,
+            params_to_args.get(f.type, f.type),
+            dataclasses.field(
+                default=f.default,
+                # ignoring mypy: https://github.com/python/mypy/issues/6910
+                default_factory=f.default_factory,  # type: ignore
+                init=f.init,
+                metadata=f.metadata,
+            ),
+        )
+        for f in dataclasses.fields(base_dataclass)
+    ]
+    non_generic_dataclass = dataclasses.make_dataclass(
+        cls_name=f"{base_dataclass.__name__}{type_arguments}", fields=non_generic_fields
+    )
+    return base_dataclass.__name__, dataclasses.fields(non_generic_dataclass)
+
+
 def NewType(
     name: str,
     typ: Type[_U],

tests/test_class_schema.py

@@ -457,6 +457,33 @@ class Meta:
         self.assertNotIn("no_init", class_schema(NoInit)().fields)
         self.assertIn("no_init", class_schema(Init)().fields)
 
+    def test_generic_dataclass(self):
+        T = typing.TypeVar("T")
+
+        @dataclasses.dataclass
+        class SimpleGeneric(typing.Generic[T]):
+            data: T
+
+        @dataclasses.dataclass
+        class Nested:
+            data: SimpleGeneric[int]
+
+        schema_s = class_schema(SimpleGeneric[str])()
+        self.assertEqual(SimpleGeneric(data="a"), schema_s.load({"data": "a"}))
+        self.assertEqual(schema_s.dump(SimpleGeneric(data="a")), {"data": "a"})
+        with self.assertRaises(ValidationError):
+            schema_s.load({"data": 2})
+
+        schema_n = class_schema(Nested)()
+        self.assertEqual(
+            Nested(data=SimpleGeneric(1)), schema_n.load({"data": {"data": 1}})
+        )
+        self.assertEqual(
+            schema_n.dump(Nested(data=SimpleGeneric(data=1))), {"data": {"data": 1}}
+        )
+        with self.assertRaises(ValidationError):
+            schema_n.load({"data": {"data": "str"}})
+
 
 if __name__ == "__main__":
     unittest.main()