Skip to content

Commit 0111e36

Browse files
syastrovsyastrov
committed
asdict: Support subclasses of list/dict (test missing for dict).
1 parent c26d6d7 commit 0111e36

2 files changed

Lines changed: 65 additions & 12 deletions

File tree

‎mypy/plugins/dataclasses.py‎

Lines changed: 30 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
"""Plugin that provides support for dataclasses."""
22

33
from collections import OrderedDict
4-
from typing import Dict, List, Set, Tuple, Optional, FrozenSet
4+
from typing import Dict, List, Set, Tuple, Optional, FrozenSet, Callable
55

66
from typing_extensions import Final
77

8+
from mypy.maptype import map_instance_to_supertype
89
from mypy.nodes import (
910
ARG_OPT, ARG_POS, MDEF, Argument, AssignmentStmt, CallExpr,
1011
Context, Expression, JsonDict, NameExpr, RefExpr,
@@ -416,6 +417,22 @@ def asdict_callback(ctx: FunctionContext) -> Type:
416417
return ctx.default_return_type
417418

418419

420+
def _transform_type_args_used_in_supertype(*, typ: Instance, supertype: TypeInfo, transform: Callable[[Instance], Type]) -> \
421+
List[Type]:
422+
"""For each type arg used in typ, call transform function if the arg is used in the given supertype."""
423+
supertype_instance = map_instance_to_supertype(typ, supertype)
424+
supertype_args = set([arg for arg in supertype_instance.args if isinstance(arg, Instance)])
425+
new_args = []
426+
# Transform existing type args if they are part of the supertype's type args
427+
for arg in typ.args:
428+
if isinstance(arg, Instance) and arg in supertype_args:
429+
new_arg = transform(arg)
430+
else:
431+
new_arg = arg
432+
new_args.append(new_arg)
433+
return new_args
434+
435+
419436
def _type_asdict(api: CheckerPluginInterface, context: Context, typ: Type) -> Type:
420437
"""Convert dataclasses into TypedDicts, recursively looking into built-in containers.
421438
@@ -447,15 +464,19 @@ def _type_asdict_inner(typ: Type, seen_dataclasses: FrozenSet[str]) -> Type:
447464
fields[attr.name] = _type_asdict_inner(typ, seen_dataclasses)
448465
return make_anonymous_typeddict(api, fields=fields, required_keys=set(fields.keys()))
449466
elif info.has_base('builtins.list'):
450-
# TODO: Support subclasses properly
451-
assert len(typ.args) == 1
452-
arg = typ.args[0]
453-
return Instance(typ.type, [_type_asdict_inner(arg, seen_dataclasses)])
467+
new_args = _transform_type_args_used_in_supertype(
468+
typ=typ,
469+
supertype=api.named_generic_type('builtins.list', []).type,
470+
transform=lambda arg: _type_asdict_inner(arg, seen_dataclasses)
471+
)
472+
return Instance(typ.type, new_args)
454473
elif info.has_base('builtins.dict'):
455-
# TODO: Support subclasses properly
456-
assert len(typ.args) == 2
457-
return Instance(typ.type, [_type_asdict_inner(typ.args[0], seen_dataclasses),
458-
_type_asdict_inner(typ.args[1], seen_dataclasses)])
474+
new_args = _transform_type_args_used_in_supertype(
475+
typ=typ,
476+
supertype=api.named_generic_type('builtins.dict', []).type,
477+
transform=lambda arg: _type_asdict_inner(arg, seen_dataclasses)
478+
)
479+
return Instance(typ.type, new_args)
459480
elif isinstance(typ, TupleType):
460481
# TODO: Support subclasses/namedtuples properly
461482
return TupleType([_type_asdict_inner(item, seen_dataclasses) for item in typ.items],

‎test-data/unit/check-dataclasses.test‎

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1130,28 +1130,60 @@ reveal_type(asdict(Person('John', 32), dict_factory=my_dict_factory)) # N: Reve
11301130

11311131
[case testDataclassesAsdictList]
11321132
from dataclasses import dataclass, asdict
1133-
from typing import List, Any
1133+
from typing import List, Any, TypeVar, Generic
11341134

11351135
@dataclass
11361136
class Person:
11371137
name: str
11381138
age: int
11391139

1140+
1141+
_T = TypeVar("_T")
1142+
class MyList(List[_T]):
1143+
pass
1144+
1145+
_X = TypeVar("_X")
1146+
class MyListWith2TypeVars(List[_T], Generic[_T, _X]):
1147+
foo: _X
1148+
1149+
_C = TypeVar("_C", Person, int)
1150+
1151+
class MyListWithConstraint(List[_C], Generic[_C]):
1152+
pass
1153+
11401154
@dataclass
11411155
class Course:
11421156
participants: List[Person]
11431157
any_list: List[Any]
11441158
list_no_generic: list
1145-
1146-
instance = Course(participants=[Person("Joe", 32)], any_list=[], list_no_generic=[])
1159+
list_subclass: MyList[Person]
1160+
list_subclass_2_typevars: MyListWith2TypeVars[Person, int]
1161+
list_subclass_with_constraint: MyListWithConstraint[Person]
1162+
1163+
instance = Course(
1164+
participants=[Person("Joe", 32)],
1165+
any_list=[],
1166+
list_no_generic=[],
1167+
list_subclass=MyList([]),
1168+
list_subclass_2_typevars=MyListWith2TypeVars[Person, int]([Person("John", 23)]),
1169+
list_subclass_with_constraint=MyListWithConstraint([Person("Tim", 29)])
1170+
)
11471171
result = asdict(instance)
11481172
reveal_type(result['participants']) # N: Revealed type is 'builtins.list[TypedDict({'name': builtins.str, 'age': builtins.int})]'
11491173
reveal_type(result['any_list']) # N: Revealed type is 'builtins.list[Any]'
11501174
reveal_type(result['list_no_generic']) # N: Revealed type is 'builtins.list[Any]'
1175+
reveal_type(result['list_subclass']) # N: Revealed type is '__main__.MyList[TypedDict({'name': builtins.str, 'age': builtins.int})]'
1176+
reveal_type(result['list_subclass_2_typevars']) # N: Revealed type is '__main__.MyListWith2TypeVars[TypedDict({'name': builtins.str, 'age': builtins.int}), builtins.int]'
1177+
1178+
# TODO: actually, error should be expected since constraints are violated.
1179+
# So how to handle this? Should we just say it is not actually a MyListWithConstraint, but Any?
1180+
# Or should it genuinely throw a type checking error
1181+
reveal_type(result['list_subclass_with_constraint']) # N: Revealed type is '__main__.MyListWithConstraint[TypedDict({'name': builtins.str, 'age': builtins.int})]'
11511182

11521183
[typing fixtures/typing-full.pyi]
11531184
[builtins fixtures/list.pyi]
11541185

1186+
11551187
[case testDataclassesAsdictDict]
11561188
from dataclasses import dataclass, asdict
11571189
from typing import Dict

0 commit comments

Comments
 (0)