refactor: add _SchemaContext.get_type_mapping

Author: dairiki

marshmallow_dataclass/__init__.py

@@ -433,6 +433,35 @@ def class_schema(
         return _internal_class_schema(clazz, base_schema)
 
 
+class _TypeMapping:
+    """Helper for looking up field types in a chained list of TYPE_MAPPINGs"""
+
+    def __init__(self, *mappings: Mapping[Any, Type[marshmallow.fields.Field]]) -> None:
+        self.mappings = mappings
+
+    _Field = TypeVar("_Field", bound=marshmallow.fields.Field)
+
+    @overload
+    def get(self, typ: object, default: Type[_Field]) -> Type[_Field]:
+        ...
+
+    @overload
+    def get(
+        self, typ: object, default: None = None
+    ) -> Optional[Type[marshmallow.fields.Field]]:
+        ...
+
+    def get(
+        self, typ: object, default: Optional[Type[_Field]] = None
+    ) -> Optional[Type[marshmallow.fields.Field]]:
+        for mapping in self.mappings:
+            try:
+                return mapping[typ]
+            except KeyError:
+                pass
+        return default
+
+
 @dataclasses.dataclass
 class _SchemaContext:
     """Global context for an invocation of class_schema."""
@@ -442,6 +471,18 @@ class _SchemaContext:
     base_schema: Optional[Type[marshmallow.Schema]] = None
     seen_classes: Dict[type, str] = dataclasses.field(default_factory=dict)
 
+    def get_type_mapping(
+        self, include_marshmallow_default: bool = False
+    ) -> _TypeMapping:
+        default_mapping = marshmallow.Schema.TYPE_MAPPING
+        if self.base_schema is not None:
+            mappings = [self.base_schema.TYPE_MAPPING]
+            if include_marshmallow_default:
+                mappings.append(default_mapping)
+        else:
+            mappings = [default_mapping]
+        return _TypeMapping(*mappings)
+
     def __enter__(self) -> "_SchemaContext":
         _schema_ctx_stack.push(self)
         return self
@@ -534,10 +575,9 @@ def _internal_class_schema(
 
 def _field_by_type(typ: Union[type, Any]) -> Optional[Type[marshmallow.fields.Field]]:
     # FIXME: remove this function
-    base_schema = _schema_ctx_stack.top.base_schema
-    return (
-        base_schema and base_schema.TYPE_MAPPING.get(typ)
-    ) or marshmallow.Schema.TYPE_MAPPING.get(typ)
+    schema_ctx = _schema_ctx_stack.top
+    type_mapping = schema_ctx.get_type_mapping(include_marshmallow_default=True)
+    return type_mapping.get(typ)
 
 
 def _field_by_supertype(
@@ -605,15 +645,12 @@ def _field_for_generic_type(
     arguments = typing_inspect.get_args(typ, True)
     if origin:
         # Override base_schema.TYPE_MAPPING to change the class used for generic types below
-        base_schema = _schema_ctx_stack.top.base_schema
-        type_mapping = base_schema.TYPE_MAPPING if base_schema else {}
+        schema_ctx = _schema_ctx_stack.top
+        type_mapping = schema_ctx.get_type_mapping()
 
         if origin in (list, List):
             child_type = _field_for_schema(arguments[0])
-            list_type = cast(
-                Type[marshmallow.fields.List],
-                type_mapping.get(List, marshmallow.fields.List),
-            )
+            list_type = type_mapping.get(List, default=marshmallow.fields.List)
             return list_type(child_type, **metadata)
         if origin in (collections.abc.Sequence, Sequence) or (
             origin in (tuple, Tuple)
@@ -640,15 +677,10 @@ def _field_for_generic_type(
             )
         if origin in (tuple, Tuple):
             children = tuple(_field_for_schema(arg) for arg in arguments)
-            tuple_type = cast(
-                Type[marshmallow.fields.Tuple],
-                type_mapping.get(  # type:ignore[call-overload]
-                    Tuple, marshmallow.fields.Tuple
-                ),
-            )
+            tuple_type = type_mapping.get(Tuple, default=marshmallow.fields.Tuple)
             return tuple_type(children, **metadata)
         elif origin in (dict, Dict, collections.abc.Mapping, Mapping):
-            dict_type = type_mapping.get(Dict, marshmallow.fields.Dict)
+            dict_type = type_mapping.get(Dict, default=marshmallow.fields.Dict)
             return dict_type(
                 keys=_field_for_schema(arguments[0]),
                 values=_field_for_schema(arguments[1]),