support nested generic dataclasses

Author: onursatici

marshmallow_dataclass/__init__.py

@@ -34,6 +34,7 @@ class User:
       })
       Schema: ClassVar[Type[Schema]] = Schema # For the type checker
 """
+
 import collections.abc
 import dataclasses
 import inspect
@@ -43,14 +44,11 @@ class User:
 import warnings
 from enum import Enum
 from functools import lru_cache, partial
+from typing import Any, Callable, Dict, FrozenSet, List, Mapping
+from typing import NewType as typing_NewType
 from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    Mapping,
-    NewType as typing_NewType,
     Optional,
+    Sequence,
     Set,
     Tuple,
     Type,
@@ -59,16 +57,13 @@ class User:
     cast,
     get_type_hints,
     overload,
-    Sequence,
-    FrozenSet,
 )
 
 import marshmallow
 import typing_inspect
 
 from marshmallow_dataclass.lazy_class_attribute import lazy_class_attribute
 
-
 if sys.version_info >= (3, 11):
     from typing import dataclass_transform
 elif sys.version_info >= (3, 7):
@@ -105,8 +100,7 @@ def dataclass(
     frozen: bool = False,
     base_schema: Optional[Type[marshmallow.Schema]] = None,
     cls_frame: Optional[types.FrameType] = None,
-) -> Type[_U]:
-    ...
+) -> Type[_U]: ...
 
 
 @overload
@@ -119,8 +113,7 @@ def dataclass(
     frozen: bool = False,
     base_schema: Optional[Type[marshmallow.Schema]] = None,
     cls_frame: Optional[types.FrameType] = None,
-) -> Callable[[Type[_U]], Type[_U]]:
-    ...
+) -> Callable[[Type[_U]], Type[_U]]: ...
 
 
 # _cls should never be specified by keyword, so start it with an
@@ -179,24 +172,21 @@ def dataclass(
 
 
 @overload
-def add_schema(_cls: Type[_U]) -> Type[_U]:
-    ...
+def add_schema(_cls: Type[_U]) -> Type[_U]: ...
 
 
 @overload
 def add_schema(
     base_schema: Optional[Type[marshmallow.Schema]] = None,
-) -> Callable[[Type[_U]], Type[_U]]:
-    ...
+) -> Callable[[Type[_U]], Type[_U]]: ...
 
 
 @overload
 def add_schema(
     _cls: Type[_U],
     base_schema: Optional[Type[marshmallow.Schema]] = None,
     cls_frame: Optional[types.FrameType] = None,
-) -> Type[_U]:
-    ...
+) -> Type[_U]: ...
 
 
 def add_schema(_cls=None, base_schema=None, cls_frame=None):
@@ -386,20 +376,27 @@ def class_schema(
         del current_frame
     _RECURSION_GUARD.seen_classes = {}
     try:
-        return _internal_class_schema(clazz, base_schema, clazz_frame)
+        return _internal_class_schema(clazz, base_schema, clazz_frame, None)
     finally:
         _RECURSION_GUARD.seen_classes.clear()
 
 
+def _dataclass_fields(clazz: type) -> Tuple[dataclasses.Field, ...]:
+    if _is_generic_alias_of_dataclass(clazz):
+        clazz = typing_inspect.get_origin(clazz)
+    return dataclasses.fields(clazz)
+
+
 @lru_cache(maxsize=MAX_CLASS_SCHEMA_CACHE_SIZE)
 def _internal_class_schema(
     clazz: type,
     base_schema: Optional[Type[marshmallow.Schema]] = None,
     clazz_frame: Optional[types.FrameType] = None,
+    generic_params_to_args: Optional[Tuple[Tuple[type, type], ...]] = None,
 ) -> Type[marshmallow.Schema]:
     _RECURSION_GUARD.seen_classes[clazz] = clazz.__name__
     try:
-        class_name, fields = _dataclass_name_and_fields(clazz)
+        fields = _dataclass_fields(clazz)
     except TypeError:  # Not a dataclass
         try:
             warnings.warn(
@@ -414,7 +411,9 @@ def _internal_class_schema(
                 "****** WARNING ******"
             )
             created_dataclass: type = dataclasses.dataclass(clazz)
-            return _internal_class_schema(created_dataclass, base_schema, clazz_frame)
+            return _internal_class_schema(
+                created_dataclass, base_schema, clazz_frame, generic_params_to_args
+            )
         except Exception as exc:
             raise TypeError(
                 f"{getattr(clazz, '__name__', repr(clazz))} is not a dataclass and cannot be turned into one."
@@ -430,10 +429,11 @@ def _internal_class_schema(
     # Determine whether we should include non-init fields
     include_non_init = getattr(getattr(clazz, "Meta", None), "include_non_init", False)
 
+    if _is_generic_alias_of_dataclass(clazz) and generic_params_to_args is None:
+        generic_params_to_args = _generic_params_to_args(clazz)
+
+    type_hints = _dataclass_type_hints(clazz, clazz_frame, generic_params_to_args)
     # Update the schema members to contain marshmallow fields instead of dataclass fields
-    type_hints = get_type_hints(
-        clazz, localns=clazz_frame.f_locals if clazz_frame else None
-    )
     attributes.update(
         (
             field.name,
@@ -443,13 +443,14 @@ def _internal_class_schema(
                 field.metadata,
                 base_schema,
                 clazz_frame,
+                generic_params_to_args,
             ),
         )
         for field in fields
         if field.init or include_non_init
     )
 
-    schema_class = type(class_name, (_base_schema(clazz, base_schema),), attributes)
+    schema_class = type(clazz.__name__, (_base_schema(clazz, base_schema),), attributes)
     return cast(Type[marshmallow.Schema], schema_class)
 
 
@@ -584,7 +585,7 @@ def _field_for_generic_type(
                 ),
             )
             return tuple_type(children, **metadata)
-        elif origin in (dict, Dict, collections.abc.Mapping, Mapping):
+        if origin in (dict, Dict, collections.abc.Mapping, Mapping):
             dict_type = type_mapping.get(Dict, marshmallow.fields.Dict)
             return dict_type(
                 keys=field_for_schema(
@@ -636,6 +637,7 @@ def field_for_schema(
     metadata: Optional[Mapping[str, Any]] = None,
     base_schema: Optional[Type[marshmallow.Schema]] = None,
     typ_frame: Optional[types.FrameType] = None,
+    generic_params_to_args: Optional[Tuple[Tuple[type, type], ...]] = None,
 ) -> marshmallow.fields.Field:
     """
     Get a marshmallow Field corresponding to the given python type.
@@ -769,7 +771,7 @@ def field_for_schema(
         nested_schema
         or forward_reference
         or _RECURSION_GUARD.seen_classes.get(typ)
-        or _internal_class_schema(typ, base_schema, typ_frame)  # type: ignore [arg-type]
+        or _internal_class_schema(typ, base_schema, typ_frame, generic_params_to_args)  # type: ignore [arg-type]
     )
 
     return marshmallow.fields.Nested(nested, **metadata)
@@ -823,35 +825,33 @@ def _is_generic_alias_of_dataclass(clazz: type) -> bool:
     )
 
 
-# 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)
-
+def _generic_params_to_args(clazz: type) -> Tuple[Tuple[type, type], ...]:
     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)
+    return tuple(zip(base_parameters, type_arguments))
+
+
+def _dataclass_type_hints(
+    clazz: type,
+    clazz_frame: types.FrameType = None,
+    generic_params_to_args: Optional[Tuple[Tuple[type, type], ...]] = None,
+) -> Mapping[str, type]:
+    localns = clazz_frame.f_locals if clazz_frame else None
+    if not _is_generic_alias_of_dataclass(clazz):
+        return get_type_hints(clazz, localns=localns)
+    # dataclass is generic
+    generic_type_hints = get_type_hints(typing_inspect.get_origin(clazz), localns)
+    generic_params_map = dict(generic_params_to_args if generic_params_to_args else {})
+
+    def _get_hint(_t: type) -> type:
+        if isinstance(_t, TypeVar):
+            return generic_params_map[_t]
+        return _t
+
+    return {
+        field_name: _get_hint(typ) for field_name, typ in generic_type_hints.items()
+    }
 
 
 def NewType(

tests/test_class_schema.py

@@ -14,7 +14,7 @@
 from marshmallow.fields import Field, UUID as UUIDField, List as ListField, Integer
 from marshmallow.validate import Validator
 
-from marshmallow_dataclass import class_schema, NewType
+from marshmallow_dataclass import class_schema, NewType, _is_generic_alias_of_dataclass
 
 
 class TestClassSchema(unittest.TestCase):