dataclasses.asdict: Handle list/tuple/dicts/namedtuples recursively.

python · syastrov · Jan 21, 2020 · Jan 28, 2020 · Feb 8, 2020 · Feb 8, 2020
commit 05a554cd3967ccac0abb8f43c765599bab3b47a4
diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py
@@ -10,14 +10,14 @@
     Context, Expression, JsonDict, NameExpr, RefExpr,
     SymbolTableNode, TempNode, TypeInfo, Var, TypeVarExpr, PlaceholderNode
 )
-from mypy.plugin import ClassDefContext, FunctionContext
+from mypy.plugin import ClassDefContext, FunctionContext, CheckerPluginInterface
 from mypy.plugin import SemanticAnalyzerPluginInterface
 from mypy.plugins.common import add_method, _get_decorator_bool_argument, make_anonymous_typeddict
 from mypy.plugins.common import (
     deserialize_and_fixup_type,
 )
 from mypy.server.trigger import make_wildcard_trigger
-from mypy.types import Instance, NoneType, TypeVarDef, TypeVarType, get_proper_type, Type
+from mypy.types import Instance, NoneType, TypeVarDef, TypeVarType, get_proper_type, Type, TupleType
 
 # The set of decorators that generate dataclasses.
 dataclass_makers = {
@@ -402,14 +402,46 @@ def asdict_callback(ctx: FunctionContext) -> Type:
         dataclass_instance = arg_types[0]
         if isinstance(dataclass_instance, Instance):
             info = dataclass_instance.type
-            if is_type_dataclass(info):
-                attrs = info.metadata['dataclass']['attributes']
-                fields = OrderedDict()  # type: OrderedDict[str, Type]
-                for data in attrs:
-                    attr = DataclassAttribute.deserialize(info, data, ctx.api)
-                    sym_node = info.names[attr.name]
-                    typ = sym_node.type
-                    assert typ is not None
-                    fields[attr.name] = typ
-                return make_anonymous_typeddict(ctx.api, fields=fields, required_keys=set(fields.keys()))
+            if not is_type_dataclass(info):
+                ctx.api.fail('asdict() should be called on dataclass instances', dataclass_instance)
+            return _type_asdict_inner(ctx.api, dataclass_instance)
     return ctx.default_return_type
+
+
+def _type_asdict_inner(api: CheckerPluginInterface, typ: Type) -> Type:
+    """Convert dataclasses into TypedDicts, recursively looking into built-in containers.
+
+    It will look for dataclasses inside of tuples, lists, and dicts and convert them to TypedDicts.
+    """
+    # TODO: detect recursive references and replace them with Any, and
+    #  perhaps generate an error about recursive types not being supported
+    if isinstance(typ, Instance):
+        info = typ.type
+        if is_type_dataclass(info):
+            attrs = info.metadata['dataclass']['attributes']
+            fields = OrderedDict()  # type: OrderedDict[str, Type]
+            for data in attrs:
+                attr = DataclassAttribute.deserialize(info, data, api)
+                sym_node = info.names[attr.name]
+                typ = sym_node.type
+                assert typ is not None
+                fields[attr.name] = _type_asdict_inner(api, typ)
+            return make_anonymous_typeddict(api, fields=fields, required_keys=set(fields.keys()))
+        elif info.has_base('builtins.list'):
+            # TODO: Support subclasses properly
+            # TODO: Does List[Any] work?
+            assert len(typ.args) == 1
+            arg = typ.args[0]
+            return Instance(typ.type, [_type_asdict_inner(api, arg)])
+        elif info.has_base('builtins.dict'):
+            # TODO: Support subclasses properly
+            assert len(typ.args) == 2
+            return Instance(typ.type, [_type_asdict_inner(api, typ.args[0]), _type_asdict_inner(api, typ.args[1])])
+    elif isinstance(typ, TupleType):
+        # Is this really the proper way to do it? or is it an Instance instead?
+        # TODO: Handle namedtuples properly.
+        # TODO: Handle the fact that the tuple type is still an "instance" of the original tuple type
+        return TupleType([_type_asdict_inner(api, item) for item in typ.items],
+                         # TODO: Recalculate fallback using tuple_fallback?
+                         typ.partial_fallback, implicit=typ.implicit)
+    return typ
diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test
@@ -1058,43 +1058,139 @@ reveal_type(asdict(Person('John', 32)))  # N: Revealed type is 'TypedDict({'name
 [builtins fixtures/list.pyi]
 [builtins fixtures/dict.pyi]
 
-[case testDataclassesAsdictRecursive]
+
+[case testDataclassesAsdictList]
 from dataclasses import dataclass, asdict
-from typing import List, Dict, Tuple, NamedTuple
+from typing import List
 
 @dataclass
 class Person:
     name: str
     age: int
 
-class NT(NamedTuple):
-    nt: Dict[str, Tuple[int, Person]]
+@dataclass
+class Course:
+    participants: List[Person]
 
-class PersonList(List[Person]):
-    pass
+instance = Course(participants=[Person("Joe", 32)])
+result = asdict(instance)
+reveal_type(result['participants']) # N: Revealed type is 'builtins.list[TypedDict({'name': builtins.str, 'age': builtins.int})]'
 
-@dataclass
-class Recurse:
-    foo: NT
-    bar: PersonList
 
-# asdict should recurse into dataclasses, dicts, lists, and tuples
-instance = Recurse(NT({'key': (1, Person("Joe", 32))}), PersonList([Person("Dave", 32)]))
+[typing fixtures/typing-full.pyi]
+[builtins fixtures/list.pyi]
+
 
-# Problem 1: There is a problem when subclassing List (and I assume also Dict?)
-# The type returned will contain an instance of PersonList, but it will contain dicts instead of Person instances
-# How can we say it still is an instance of PersonList when it has values of the wrong type?
-# We could copy the existing PersonList type and modify its generic parameter to List...
-# but at runtime: isinstance(asdict(instance)['bar'], PersonList) == True
+[case testDataclassesAsdictDict]
+from dataclasses import dataclass, asdict
+from typing import Dict
+
+@dataclass
+class Person:
+    name: str
+    age: int
 
-# Problem 2: The same problem occurs for namedtuples: they should be instances of the original namedtuple class,
-# but their values which originally were instances of dataclasses should be dicts instead.
+@dataclass
+class Course:
+    participants_by_name: Dict[str, Person]
 
-# I'm expecting the revealed type to be something like this, but this wouldn't really be correct due to the above...
-reveal_type(asdict(instance)) # N: Revealed type is 'TypedDict({'foo': builtins.tuple[Dict[str, Tuple[int, TypedDict({'name': builtins.str, 'age': builtins.int}), fallback=__main__.NT]]], 'bar': __main__.PersonList})'
+instance = Course(participants_by_name={"Joe": Person("Joe", 32)})
+result = asdict(instance)
+reveal_type(result['participants_by_name']) # N: Revealed type is 'builtins.dict[builtins.str, TypedDict({'name': builtins.str, 'age': builtins.int})]'
 
-# TODO: Handle dict_factory param -- what type to give the return value then?
 
 [typing fixtures/typing-full.pyi]
-[builtins fixtures/list.pyi]
 [builtins fixtures/dict.pyi]
+
+[case testDataclassesAsdictTuple]
+from dataclasses import dataclass, asdict
+from typing import Tuple
+
+@dataclass
+class Person:
+    name: str
+    age: int
+
+@dataclass
+class Course:
+    partners: Tuple[Person, Person]
+
+instance = Course(partners=(Person("Joe", 32), Person("John", 23)))
+result = asdict(instance)
+reveal_type(result['partners']) # N: Revealed type is 'Tuple[TypedDict({'name': builtins.str, 'age': builtins.int}), TypedDict({'name': builtins.str, 'age': builtins.int})]'
+
+
+[typing fixtures/typing-full.pyi]
+[builtins fixtures/tuple.pyi]
+
+
+[case testDataclassesAsdictNamedTuple]
+from dataclasses import dataclass, asdict
+from typing import NamedTuple
+
+@dataclass
+class Person:
+    name: str
+    age: int
+
+
+class Staff(NamedTuple):
+    teacher: Person
+    assistant: Person
+
+@dataclass
+class Course:
+    staff: Staff
+
+instance = Course(staff=Staff(teacher=Person("Joe", 32), assistant=Person("John", 23)))
+result = asdict(instance)
+reveal_type(result['staff']) # N: Revealed type is 'Tuple[TypedDict({'name': builtins.str, 'age': builtins.int}), TypedDict({'name': builtins.str, 'age': builtins.int}), fallback=__main__.Staff]'
+staff = result['staff']
+if isinstance(staff, Staff):
+    # TODO: Make this pass: at runtime it succeeds
+    reveal_type(staff) # N: Revealed type is 'Tuple[TypedDict({'name': builtins.str, 'age': builtins.int}), TypedDict({'name': builtins.str, 'age': builtins.int}), fallback=__main__.Staff]'
+
+[typing fixtures/typing-full.pyi]
+[builtins fixtures/tuple.pyi]
+
+#
+# [case testDataclassesAsdictRecursive]
+# from dataclasses import dataclass, asdict
+# from typing import List, Dict, Tuple, NamedTuple
+#
+# @dataclass
+# class Person:
+#     name: str
+#     age: int
+#
+# class NT(NamedTuple):
+#     nt: Dict[str, Tuple[int, Person]]
+#
+# class PersonList(List[Person]):
+#     pass
+#
+# @dataclass
+# class Recurse:
+#     foo: NT
+#     bar: PersonList
+#
+# # asdict should recurse into dataclasses, dicts, lists, and tuples
+# instance = Recurse(NT({'key': (1, Person("Joe", 32))}), PersonList([Person("Dave", 32)]))
+#
+# # Problem 1: There is a problem when subclassing List (and I assume also Dict?)
+# # The type returned will contain an instance of PersonList, but it will contain dicts instead of Person instances
+# # How can we say it still is an instance of PersonList when it has values of the wrong type?
+# # We could copy the existing PersonList type and modify its generic parameter to List...
+# # but at runtime: isinstance(asdict(instance)['bar'], PersonList) == True
+#
+# # Problem 2: The same problem occurs for namedtuples: they should be instances of the original namedtuple class,
+# # but their values which originally were instances of dataclasses should be dicts instead.
+#
+# # I'm expecting the revealed type to be something like this, but this wouldn't really be correct due to the above...
+# reveal_type(asdict(instance)) # N: Revealed type is 'TypedDict({'foo': builtins.tuple[Dict[str, Tuple[int, TypedDict({'name': builtins.str, 'age': builtins.int}), fallback=__main__.NT]]], 'bar': __main__.PersonList})'
+#
+# # TODO: Handle dict_factory param -- what type to give the return value then?
+#
+# [typing fixtures/typing-full.pyi]
+# [builtins fixtures/list.pyi]
+# [builtins fixtures/dict.pyi]