Add Generic dataclasses

marshmallow_dataclass/__init__.py

@@ -44,15 +44,9 @@ class User:
 import warnings
 from enum import Enum
 from functools import lru_cache, partial
+from typing import Any, Callable, Dict, FrozenSet, Generic, List, Mapping
+from typing import NewType as typing_NewType
 from typing import (
-    Any,
-    Callable,
-    Dict,
-    FrozenSet,
-    Generic,
-    List,
-    Mapping,
-    NewType as typing_NewType,
     Optional,
     Sequence,
     Set,
@@ -69,6 +63,12 @@ class User:
 import typing_extensions
 import typing_inspect
 
+from marshmallow_dataclass.generic_resolver import (
+    UnboundTypeVarError,
+    get_generic_dataclass_fields,
+    is_generic_alias,
+    is_generic_type,
+)
 from marshmallow_dataclass.lazy_class_attribute import lazy_class_attribute
 
 if sys.version_info >= (3, 9):
@@ -139,6 +139,18 @@ def _maybe_get_callers_frame(
         del frame
 
 
+def _check_decorated_type(cls: object) -> None:
+    if not isinstance(cls, type):
+        raise TypeError(f"expected a class not {cls!r}")
+    if is_generic_alias(cls):
+        # A .Schema attribute doesn't make sense on a generic alias — there's
+        # no way for it to know the generic parameters at run time.
+        raise TypeError(
+            "decorator does not support generic aliasses "
+            "(hint: use class_schema directly instead)"
+        )
+
+
 @overload
 def dataclass(
     _cls: Type[_U],
@@ -214,12 +226,18 @@ def dataclass(
     )
 
     def decorator(cls: Type[_U], stacklevel: int = 1) -> Type[_U]:
+        if cls is not None:
+            _check_decorated_type(cls)
+
         return add_schema(
             dc(cls), base_schema, cls_frame=cls_frame, stacklevel=stacklevel + 1
         )
 
     if _cls is None:
         return decorator
+
+    if _cls is not None:
+        _check_decorated_type(_cls)
     return decorator(_cls, stacklevel=stacklevel + 1)
 
 
@@ -268,6 +286,8 @@ def add_schema(_cls=None, base_schema=None, cls_frame=None, stacklevel=1):
     """
 
     def decorator(clazz: Type[_U], stacklevel: int = stacklevel) -> Type[_U]:
+        _check_decorated_type(clazz)
+
         if cls_frame is not None:
             frame = cls_frame
         else:
@@ -453,7 +473,7 @@ 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 localns is None:
         if clazz_frame is None:
@@ -518,13 +538,15 @@ def _internal_class_schema(
         # https://github.com/python/cpython/blob/3.10/Lib/typing.py#L977
         class_name = clazz._name or clazz.__origin__.__name__  # type: ignore[attr-defined]
     else:
-        class_name = clazz.__name__
+        # generic aliases do not have a __name__ prior python 3.10
+        class_name = getattr(clazz, "__name__", repr(clazz))
 
     schema_ctx.seen_classes[clazz] = class_name
 
     try:
-        # noinspection PyDataclass
-        fields: Tuple[dataclasses.Field, ...] = dataclasses.fields(clazz)
+        fields = _dataclass_fields(clazz)
+    except UnboundTypeVarError:
+        raise
     except TypeError:  # Not a dataclass
         try:
             warnings.warn(
@@ -540,6 +562,8 @@ def _internal_class_schema(
             )
             created_dataclass: type = dataclasses.dataclass(clazz)
             return _internal_class_schema(created_dataclass, base_schema)
+        except UnboundTypeVarError:
+            raise
         except Exception as exc:
             raise TypeError(
                 f"{getattr(clazz, '__name__', repr(clazz))} is not a dataclass and cannot be turned into one."
@@ -556,23 +580,19 @@ def _internal_class_schema(
     include_non_init = getattr(getattr(clazz, "Meta", None), "include_non_init", False)
 
     # Update the schema members to contain marshmallow fields instead of dataclass fields
+    type_hints = {}
+    if not is_generic_type(clazz):
+        type_hints = _get_type_hints(clazz, schema_ctx)
 
-    if sys.version_info >= (3, 9):
-        type_hints = get_type_hints(
-            clazz,
-            globalns=schema_ctx.globalns,
-            localns=schema_ctx.localns,
-            include_extras=True,
-        )
-    else:
-        type_hints = get_type_hints(
-            clazz, globalns=schema_ctx.globalns, localns=schema_ctx.localns
-        )
     attributes.update(
         (
             field.name,
             _field_for_schema(
-                type_hints[field.name],
+                (
+                    type_hints[field.name]
+                    if not is_generic_type(clazz)
+                    else _get_generic_type_hints(field.type, schema_ctx)
+                ),
                 _get_field_default(field),
                 field.metadata,
                 base_schema,
@@ -582,7 +602,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)
 
 
@@ -705,7 +725,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(arguments[0], base_schema=base_schema),
@@ -729,8 +749,12 @@ def _field_for_annotated_type(
         marshmallow_annotations = [
             arg
             for arg in arguments[1:]
-            if (inspect.isclass(arg) and issubclass(arg, marshmallow.fields.Field))
-            or isinstance(arg, marshmallow.fields.Field)
+            if _is_marshmallow_field(arg)
+            # Support `CustomGenericField[mf.String]`
+            or (
+                is_generic_type(arg)
+                and _is_marshmallow_field(typing_extensions.get_origin(arg))
+            )
         ]
         if marshmallow_annotations:
             if len(marshmallow_annotations) > 1:
@@ -838,6 +862,9 @@ def _field_for_schema(
 
     """
 
+    if isinstance(typ, TypeVar):
+        raise UnboundTypeVarError(f"can not resolve type variable {typ.__name__}")
+
     metadata = {} if metadata is None else dict(metadata)
 
     if default is not marshmallow.missing:
@@ -867,7 +894,7 @@ def _field_for_schema(
 
     # i.e.: Literal['abc']
     if typing_inspect.is_literal_type(typ):
-        arguments = typing_inspect.get_args(typ)
+        arguments = typing_extensions.get_args(typ)
         return marshmallow.fields.Raw(
             validate=(
                 marshmallow.validate.Equal(arguments[0])
@@ -879,7 +906,7 @@ def _field_for_schema(
 
     # i.e.: Final[str] = 'abc'
     if typing_inspect.is_final_type(typ):
-        arguments = typing_inspect.get_args(typ)
+        arguments = typing_extensions.get_args(typ)
         if arguments:
             subtyp = arguments[0]
         elif default is not marshmallow.missing:
@@ -952,7 +979,7 @@ def _field_for_schema(
         nested_schema
         or forward_reference
         or _schema_ctx_stack.top.seen_classes.get(typ)
-        or _internal_class_schema(typ, base_schema)  # type: ignore[arg-type] # FIXME
+        or _internal_class_schema(typ, base_schema)  # type: ignore [arg-type]
     )
 
     return marshmallow.fields.Nested(nested, **metadata)
@@ -996,6 +1023,66 @@ 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
+    """
+    is_generic = is_generic_type(clazz)
+    type_arguments = typing_extensions.get_args(clazz)
+    origin_class = typing_extensions.get_origin(clazz)
+    return (
+        is_generic
+        and len(type_arguments) > 0
+        and dataclasses.is_dataclass(origin_class)
+    )
+
+
+def _get_type_hints(
+    obj,
+    schema_ctx: _SchemaContext,
+):
+    if sys.version_info >= (3, 9):
+        type_hints = get_type_hints(
+            obj,
+            globalns=schema_ctx.globalns,
+            localns=schema_ctx.localns,
+            include_extras=True,
+        )
+    else:
+        type_hints = get_type_hints(
+            obj, globalns=schema_ctx.globalns, localns=schema_ctx.localns
+        )
+
+    return type_hints
+
+
+def _get_generic_type_hints(
+    obj,
+    schema_ctx: _SchemaContext,
+) -> type:
+    """typing.get_type_hints doesn't work with generic aliasses. But this 'hack' works."""
+
+    class X:
+        x: obj  # type: ignore[name-defined]
+
+    return _get_type_hints(X, schema_ctx)["x"]
+
+
+def _dataclass_fields(clazz: type) -> Tuple[dataclasses.Field, ...]:
+    if not is_generic_type(clazz):
+        return dataclasses.fields(clazz)
+
+    else:
+        return get_generic_dataclass_fields(clazz)
+
+
+def _is_marshmallow_field(obj) -> bool:
+    return (
+        inspect.isclass(obj) and issubclass(obj, marshmallow.fields.Field)
+    ) or isinstance(obj, marshmallow.fields.Field)
+
+
 def NewType(
     name: str,
     typ: Type[_U],