Clean up

* Differentiate between Generics and container type functions
* Tie get_args and get_origin functions to Annotated import.
* Rename function and add test to clarify forward ref use case

Author: mvanderlee

marshmallow_dataclass/__init__.py

@@ -66,7 +66,6 @@ class User:
 )
 
 import marshmallow
-import typing_extensions
 import typing_inspect
 
 from marshmallow_dataclass.generic_resolver import (
@@ -78,9 +77,9 @@ class User:
 from marshmallow_dataclass.lazy_class_attribute import lazy_class_attribute
 
 if sys.version_info >= (3, 9):
-    from typing import Annotated
+    from typing import Annotated, get_args, get_origin
 else:
-    from typing_extensions import Annotated
+    from typing_extensions import Annotated, get_args, get_origin
 
 if sys.version_info >= (3, 11):
     from typing import dataclass_transform
@@ -540,7 +539,7 @@ def _internal_class_schema(
 ) -> Type[marshmallow.Schema]:
     schema_ctx = _schema_ctx_stack.top
 
-    if typing_extensions.get_origin(clazz) is Annotated and sys.version_info < (3, 10):
+    if get_origin(clazz) is Annotated and sys.version_info < (3, 10):
         # 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:
@@ -597,7 +596,7 @@ def _internal_class_schema(
                 (
                     type_hints[field.name]
                     if not is_generic_type(clazz)
-                    else _get_type_hint_of_generic_object(field.type, schema_ctx)
+                    else _resolve_forward_type_refs(field.type, schema_ctx)
                 ),
                 _get_field_default(field),
                 field.metadata,
@@ -659,8 +658,8 @@ def _field_by_supertype(
         )
 
 
-def _generic_type_add_any(typ: type) -> type:
-    """if typ is generic type without arguments, replace them by Any."""
+def _container_type_add_any(typ: type) -> type:
+    """if typ is container type without arguments, replace them by Any."""
     if typ is list or typ is List:
         typ = List[Any]
     elif typ is dict or typ is Dict:
@@ -676,18 +675,20 @@ def _generic_type_add_any(typ: type) -> type:
     return typ
 
 
-def _field_for_generic_type(
+def _field_for_container_type(
     typ: type,
     base_schema: Optional[Type[marshmallow.Schema]],
     **metadata: Any,
 ) -> Optional[marshmallow.fields.Field]:
     """
-    If the type is a generic interface, resolve the arguments and construct the appropriate Field.
+    If the type is a container interface, resolve the arguments and construct the appropriate Field.
+
+    We use the term 'container' to differentiate from the Generic support
     """
-    origin = typing_extensions.get_origin(typ)
-    arguments = typing_extensions.get_args(typ)
+    origin = get_origin(typ)
+    arguments = get_args(typ)
     if origin:
-        # Override base_schema.TYPE_MAPPING to change the class used for generic types below
+        # Override base_schema.TYPE_MAPPING to change the class used for container types below
         type_mapping = base_schema.TYPE_MAPPING if base_schema else {}
 
         if origin in (list, List):
@@ -749,18 +750,15 @@ def _field_for_annotated_type(
     """
     If the type is an Annotated interface, resolve the arguments and construct the appropriate Field.
     """
-    origin = typing_extensions.get_origin(typ)
-    arguments = typing_extensions.get_args(typ)
+    origin = get_origin(typ)
+    arguments = get_args(typ)
     if origin and origin is Annotated:
         marshmallow_annotations = [
             arg
             for arg in arguments[1:]
             if _is_marshmallow_field(arg)
             # Support `CustomGenericField[mf.String]`
-            or (
-                is_generic_type(arg)
-                and _is_marshmallow_field(typing_extensions.get_origin(arg))
-            )
+            or (is_generic_type(arg) and _is_marshmallow_field(get_origin(arg)))
         ]
         if marshmallow_annotations:
             if len(marshmallow_annotations) > 1:
@@ -782,7 +780,7 @@ def _field_for_union_type(
     base_schema: Optional[Type[marshmallow.Schema]],
     **metadata: Any,
 ) -> Optional[marshmallow.fields.Field]:
-    arguments = typing_extensions.get_args(typ)
+    arguments = get_args(typ)
     if typing_inspect.is_union_type(typ):
         if typing_inspect.is_optional_type(typ):
             metadata["allow_none"] = metadata.get("allow_none", True)
@@ -886,8 +884,8 @@ def _field_for_schema(
     if predefined_field:
         return predefined_field
 
-    # Generic types specified without type arguments
-    typ = _generic_type_add_any(typ)
+    # Container types (generics like List) specified without type arguments
+    typ = _container_type_add_any(typ)
 
     # Base types
     field = _field_by_type(typ, base_schema)
@@ -900,7 +898,7 @@ def _field_for_schema(
 
     # i.e.: Literal['abc']
     if typing_inspect.is_literal_type(typ):
-        arguments = typing_extensions.get_args(typ)
+        arguments = get_args(typ)
         return marshmallow.fields.Raw(
             validate=(
                 marshmallow.validate.Equal(arguments[0])
@@ -912,7 +910,7 @@ def _field_for_schema(
 
     # i.e.: Final[str] = 'abc'
     if typing_inspect.is_final_type(typ):
-        arguments = typing_extensions.get_args(typ)
+        arguments = get_args(typ)
         if arguments:
             subtyp = arguments[0]
         elif default is not marshmallow.missing:
@@ -953,10 +951,10 @@ def _field_for_schema(
     if union_field:
         return union_field
 
-    # Generic types
-    generic_field = _field_for_generic_type(typ, base_schema, **metadata)
-    if generic_field:
-        return generic_field
+    # Container types
+    container_field = _field_for_container_type(typ, base_schema, **metadata)
+    if container_field:
+        return container_field
 
     # typing.NewType returns a function (in python <= 3.9) or a class (python >= 3.10) with a
     # __supertype__ attribute
@@ -1034,9 +1032,7 @@ 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 is_generic_alias(clazz) and dataclasses.is_dataclass(
-        typing_extensions.get_origin(clazz)
-    )
+    return is_generic_alias(clazz) and dataclasses.is_dataclass(get_origin(clazz))
 
 
 def _get_type_hints(
@@ -1058,11 +1054,13 @@ def _get_type_hints(
     return type_hints
 
 
-def _get_type_hint_of_generic_object(
+def _resolve_forward_type_refs(
     obj,
     schema_ctx: _SchemaContext,
 ) -> type:
-    """typing.get_type_hints doesn't work with generic aliases, i.e.: A[int]. But this 'hack' works."""
+    """
+    Resolve forward references, mainly applies to Generics i.e.: `A["int"]` -> `A[int]`
+    """
 
     class X:
         x: obj  # type: ignore[name-defined]

marshmallow_dataclass/generic_resolver.py

@@ -11,13 +11,12 @@
     TypeVar,
 )
 
-import typing_extensions
 import typing_inspect
 
 if sys.version_info >= (3, 9):
-    from typing import Annotated
+    from typing import Annotated, get_args, get_origin
 else:
-    from typing_extensions import Annotated
+    from typing_extensions import Annotated, get_args, get_origin
 
 _U = TypeVar("_U")
 
@@ -81,15 +80,15 @@ class A(Generic[T]):
     ``A[int]`` is a _generic alias_ (while ``A`` is a *generic type*, but not a *generic alias*).
     """
     is_generic = is_generic_type(clazz)
-    type_arguments = typing_extensions.get_args(clazz)
+    type_arguments = get_args(clazz)
     return is_generic and len(type_arguments) > 0
 
 
 def is_generic_type(clazz: type) -> bool:
     """
     typing_inspect.is_generic_type explicitly ignores Union and Tuple
     """
-    origin = typing_extensions.get_origin(clazz)
+    origin = get_origin(clazz)
     return origin is not Annotated and (
         (isinstance(clazz, type) and issubclass(clazz, Generic))  # type: ignore[arg-type]
         or isinstance(clazz, typing_inspect.typingGenericAlias)
@@ -108,7 +107,7 @@ def _resolve_typevars(clazz: type) -> Dict[type, Dict[TypeVar, _Future]]:
     # Loop in reversed order and iteratively resolve types
     for subclass in reversed(clazz.mro()):
         if issubclass(subclass, Generic) and hasattr(subclass, "__orig_bases__"):  # type: ignore[arg-type]
-            args = typing_extensions.get_args(subclass.__orig_bases__[0])
+            args = get_args(subclass.__orig_bases__[0])
 
             if parent_class and args_by_class.get(parent_class):
                 subclass_generic_params_to_args: List[Tuple[TypeVar, _Future]] = []
@@ -129,8 +128,8 @@ def _resolve_typevars(clazz: type) -> Dict[type, Dict[TypeVar, _Future]]:
 
     # clazz itself is a generic alias i.e.: A[int]. So it hold the last types.
     if is_generic_alias(clazz):
-        origin = typing_extensions.get_origin(clazz)
-        args = typing_extensions.get_args(clazz)
+        origin = get_origin(clazz)
+        args = get_args(clazz)
         for (_arg, future), potential_type in zip(args_by_class[origin], args):  # type: ignore[index]
             if not isinstance(potential_type, TypeVar):
                 future.set_result(potential_type)
@@ -154,7 +153,7 @@ def _replace_typevars(
                     resolved_generics[arg].result() if arg in resolved_generics else arg
                 )
             )
-            for arg in typing_extensions.get_args(clazz)
+            for arg in get_args(clazz)
         )
     )
 

tests/test_generics.py

@@ -315,6 +315,19 @@ class AnnotatedValue:
         with self.assertRaises(marshmallow.exceptions.ValidationError):
             schema.load({"emails": "notavalidemail"})
 
+    def test_generic_dataclass_with_forwardref(self):
+        T = typing.TypeVar("T")
+
+        @dataclasses.dataclass
+        class SimpleGeneric(typing.Generic[T]):
+            data: T
+
+        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})
+
 
 if __name__ == "__main__":
     unittest.main()