asdict: Support subclasses of list/dict (test missing for dict).

python · syastrov · Jan 21, 2020 · Jan 28, 2020 · Feb 8, 2020 · Feb 8, 2020
commit 0111e3646f174246bebedc979f5785d89ef73c05
diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py
@@ -1,10 +1,11 @@
 """Plugin that provides support for dataclasses."""
 
 from collections import OrderedDict
-from typing import Dict, List, Set, Tuple, Optional, FrozenSet
+from typing import Dict, List, Set, Tuple, Optional, FrozenSet, Callable
 
 from typing_extensions import Final
 
+from mypy.maptype import map_instance_to_supertype
 from mypy.nodes import (
     ARG_OPT, ARG_POS, MDEF, Argument, AssignmentStmt, CallExpr,
     Context, Expression, JsonDict, NameExpr, RefExpr,
@@ -416,6 +417,22 @@ def asdict_callback(ctx: FunctionContext) -> Type:
     return ctx.default_return_type
 
 
+def _transform_type_args_used_in_supertype(*, typ: Instance, supertype: TypeInfo, transform: Callable[[Instance], Type]) -> \
+        List[Type]:
+    """For each type arg used in typ, call transform function if the arg is used in the given supertype."""
+    supertype_instance = map_instance_to_supertype(typ, supertype)
+    supertype_args = set([arg for arg in supertype_instance.args if isinstance(arg, Instance)])
+    new_args = []
+    # Transform existing type args if they are part of the supertype's type args
+    for arg in typ.args:
+        if isinstance(arg, Instance) and arg in supertype_args:
+            new_arg = transform(arg)
+        else:
+            new_arg = arg
+        new_args.append(new_arg)
+    return new_args
+
+
 def _type_asdict(api: CheckerPluginInterface, context: Context, typ: Type) -> Type:
     """Convert dataclasses into TypedDicts, recursively looking into built-in containers.
 
@@ -447,15 +464,19 @@ def _type_asdict_inner(typ: Type, seen_dataclasses: FrozenSet[str]) -> Type:
                     fields[attr.name] = _type_asdict_inner(typ, seen_dataclasses)
                 return make_anonymous_typeddict(api, fields=fields, required_keys=set(fields.keys()))
             elif info.has_base('builtins.list'):
-                # TODO: Support subclasses properly
-                assert len(typ.args) == 1
-                arg = typ.args[0]
-                return Instance(typ.type, [_type_asdict_inner(arg, seen_dataclasses)])
+                new_args = _transform_type_args_used_in_supertype(
+                    typ=typ,
+                    supertype=api.named_generic_type('builtins.list', []).type,
+                    transform=lambda arg: _type_asdict_inner(arg, seen_dataclasses)
+                )
+                return Instance(typ.type, new_args)
             elif info.has_base('builtins.dict'):
-                # TODO: Support subclasses properly
-                assert len(typ.args) == 2
-                return Instance(typ.type, [_type_asdict_inner(typ.args[0], seen_dataclasses),
-                                           _type_asdict_inner(typ.args[1], seen_dataclasses)])
+                new_args = _transform_type_args_used_in_supertype(
+                    typ=typ,
+                    supertype=api.named_generic_type('builtins.dict', []).type,
+                    transform=lambda arg: _type_asdict_inner(arg, seen_dataclasses)
+                )
+                return Instance(typ.type, new_args)
         elif isinstance(typ, TupleType):
             # TODO: Support subclasses/namedtuples properly
             return TupleType([_type_asdict_inner(item, seen_dataclasses) for item in typ.items],

diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test
@@ -1130,28 +1130,60 @@ reveal_type(asdict(Person('John', 32), dict_factory=my_dict_factory))  # N: Reve
 
 [case testDataclassesAsdictList]
 from dataclasses import dataclass, asdict
-from typing import List, Any
+from typing import List, Any, TypeVar, Generic
 
 @dataclass
 class Person:
     name: str
     age: int
 
+
+_T = TypeVar("_T")
+class MyList(List[_T]):
+    pass
+
+_X = TypeVar("_X")
+class MyListWith2TypeVars(List[_T], Generic[_T, _X]):
+    foo: _X
+
+_C = TypeVar("_C", Person, int)
+
+class MyListWithConstraint(List[_C], Generic[_C]):
+    pass
+
 @dataclass
 class Course:
     participants: List[Person]
     any_list: List[Any]
     list_no_generic: list
-
-instance = Course(participants=[Person("Joe", 32)], any_list=[], list_no_generic=[])
+    list_subclass: MyList[Person]
+    list_subclass_2_typevars: MyListWith2TypeVars[Person, int]
+    list_subclass_with_constraint: MyListWithConstraint[Person]
+
+instance = Course(
+    participants=[Person("Joe", 32)],
+    any_list=[],
+    list_no_generic=[],
+    list_subclass=MyList([]),
+    list_subclass_2_typevars=MyListWith2TypeVars[Person, int]([Person("John", 23)]),
+    list_subclass_with_constraint=MyListWithConstraint([Person("Tim", 29)])
+)
 result = asdict(instance)
 reveal_type(result['participants']) # N: Revealed type is 'builtins.list[TypedDict({'name': builtins.str, 'age': builtins.int})]'
 reveal_type(result['any_list']) # N: Revealed type is 'builtins.list[Any]'
 reveal_type(result['list_no_generic']) # N: Revealed type is 'builtins.list[Any]'
+reveal_type(result['list_subclass']) # N: Revealed type is '__main__.MyList[TypedDict({'name': builtins.str, 'age': builtins.int})]'
+reveal_type(result['list_subclass_2_typevars']) # N: Revealed type is '__main__.MyListWith2TypeVars[TypedDict({'name': builtins.str, 'age': builtins.int}), builtins.int]'
+
+# TODO: actually, error should be expected since constraints are violated.
+#  So how to handle this? Should we just say it is not actually a MyListWithConstraint, but Any?
+#  Or should it genuinely throw a type checking error
+reveal_type(result['list_subclass_with_constraint']) # N: Revealed type is '__main__.MyListWithConstraint[TypedDict({'name': builtins.str, 'age': builtins.int})]'
 
 [typing fixtures/typing-full.pyi]
 [builtins fixtures/list.pyi]
 
+
 [case testDataclassesAsdictDict]
 from dataclasses import dataclass, asdict
 from typing import Dict