diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py index 66105f96a3eb5..b77633c6c602d 100644 --- a/mypy/plugins/common.py +++ b/mypy/plugins/common.py @@ -1,14 +1,14 @@ -from typing import List, Optional, Union +from typing import List, Optional, Union, Set from mypy.nodes import ( ARG_POS, MDEF, Argument, Block, CallExpr, Expression, SYMBOL_FUNCBASE_TYPES, FuncDef, PassStmt, RefExpr, SymbolTableNode, Var, JsonDict, ) -from mypy.plugin import ClassDefContext, SemanticAnalyzerPluginInterface +from mypy.plugin import ClassDefContext, SemanticAnalyzerPluginInterface, CheckerPluginInterface from mypy.semanal import set_callable_name from mypy.types import ( CallableType, Overloaded, Type, TypeVarDef, deserialize_type, get_proper_type, -) + TypedDictType, Instance, TPDICT_FB_NAMES) from mypy.typevars import fill_typevars from mypy.util import get_unique_redefinition_name from mypy.typeops import try_getting_str_literals # noqa: F401 # Part of public API @@ -134,8 +134,24 @@ def add_method( def deserialize_and_fixup_type( - data: Union[str, JsonDict], api: SemanticAnalyzerPluginInterface + data: Union[str, JsonDict], api: SemanticAnalyzerPluginInterface ) -> Type: typ = deserialize_type(data) typ.accept(TypeFixer(api.modules, allow_missing=False)) return typ + + +def get_anonymous_typeddict_type(api: CheckerPluginInterface) -> Instance: + for type_fullname in TPDICT_FB_NAMES: + try: + anonymous_typeddict_type = api.named_generic_type(type_fullname, []) + if anonymous_typeddict_type is not None: + return anonymous_typeddict_type + except KeyError: + continue + raise RuntimeError("No TypedDict fallback type found") + + +def make_anonymous_typeddict(api: CheckerPluginInterface, fields: 'OrderedDict[str, Type]', + required_keys: Set[str]) -> TypedDictType: + return TypedDictType(fields, required_keys=required_keys, fallback=get_anonymous_typeddict_type(api)) diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index 81f50d1a4c4d8..a8069380e3f4b 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -1,19 +1,26 @@ """Plugin that provides support for dataclasses.""" -from typing import Dict, List, Set, Tuple, Optional +from collections import OrderedDict +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, SymbolTableNode, TempNode, TypeInfo, Var, TypeVarExpr, PlaceholderNode ) -from mypy.plugin import ClassDefContext, SemanticAnalyzerPluginInterface +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 ( - add_method, _get_decorator_bool_argument, deserialize_and_fixup_type, + deserialize_and_fixup_type, ) -from mypy.types import Type, Instance, NoneType, TypeVarDef, TypeVarType, get_proper_type from mypy.server.trigger import make_wildcard_trigger +from mypy.typeops import tuple_fallback +from mypy.types import Instance, NoneType, TypeVarDef, TypeVarType, get_proper_type, Type, TupleType, UnionType, \ + AnyType, TypeOfAny # The set of decorators that generate dataclasses. dataclass_makers = { @@ -24,6 +31,10 @@ SELF_TVAR_NAME = '_DT' # type: Final +def is_type_dataclass(info: TypeInfo) -> bool: + return 'dataclass' in info.metadata + + class DataclassAttribute: def __init__( self, @@ -297,7 +308,7 @@ def collect_attributes(self) -> Optional[List[DataclassAttribute]]: # we'll have unmodified attrs laying around. all_attrs = attrs.copy() for info in cls.info.mro[1:-1]: - if 'dataclass' not in info.metadata: + if not is_type_dataclass(info): continue super_attrs = [] @@ -386,3 +397,80 @@ def _collect_field_args(expr: Expression) -> Tuple[bool, Dict[str, Expression]]: args[name] = arg return True, args return False, {} + + +def asdict_callback(ctx: FunctionContext) -> Type: + positional_arg_types = ctx.arg_types[0] + + if positional_arg_types: + if len(ctx.arg_types) == 2: + # We can't infer a more precise for calls where dict_factory is set. + # At least for now, typeshed stubs for asdict don't allow you to pass in `dict` as dict_factory, + # so we can't special-case that. + return ctx.default_return_type + dataclass_instance = positional_arg_types[0] + if isinstance(dataclass_instance, Instance): + info = dataclass_instance.type + if not is_type_dataclass(info): + ctx.api.fail('asdict() should be called on dataclass instances', dataclass_instance) + return _type_asdict(ctx.api, ctx.context, dataclass_instance) + return ctx.default_return_type + + +def _transform_type_args(*, typ: Instance, transform: Callable[[Instance], Type]) -> \ + List[Type]: + """For each type arg used in the Instance, call transform function on it if the arg is an Instance.""" + return [transform(arg) if isinstance(arg, Instance) else arg for arg in typ.args] + + +def _type_asdict(api: CheckerPluginInterface, context: Context, 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. + """ + + def _type_asdict_inner(typ: Type, seen_dataclasses: FrozenSet[str]) -> Type: + if isinstance(typ, UnionType): + return UnionType([_type_asdict_inner(item, seen_dataclasses) for item in typ.items]) + if isinstance(typ, Instance): + info = typ.type + if is_type_dataclass(info): + if info.fullname in seen_dataclasses: + api.fail("Recursive types are not supported in call to asdict, so falling back to Dict[str, Any]", + context) + # Note: Would be nicer to fallback to default_return_type, but that is Any (due to overloads?) + return api.named_generic_type('builtins.dict', [api.named_generic_type('builtins.str', []), + AnyType(TypeOfAny.implementation_artifact)]) + seen_dataclasses |= {info.fullname} + attrs = info.metadata['dataclass']['attributes'] + fields = OrderedDict() # type: OrderedDict[str, Type] + for data in attrs: + # TODO: DataclassAttribute.deserialize takes SemanticAnalyzerPluginInterface but we have + # CheckerPluginInterface here. + 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(typ, seen_dataclasses) + return make_anonymous_typeddict(api, fields=fields, required_keys=set(fields.keys())) + elif info.has_base('builtins.list'): + supertype_instance = map_instance_to_supertype(typ, api.named_generic_type('builtins.list', []).type) + new_args = _transform_type_args( + typ=supertype_instance, + transform=lambda arg: _type_asdict_inner(arg, seen_dataclasses) + ) + return api.named_generic_type('builtins.list', new_args) + elif info.has_base('builtins.dict'): + supertype_instance = map_instance_to_supertype(typ, api.named_generic_type('builtins.dict', []).type) + new_args = _transform_type_args( + typ=supertype_instance, + transform=lambda arg: _type_asdict_inner(arg, seen_dataclasses) + ) + return api.named_generic_type('builtins.dict', new_args) + elif isinstance(typ, TupleType): + # TODO: Support subclasses/namedtuples properly + return TupleType([_type_asdict_inner(item, seen_dataclasses) for item in typ.items], + tuple_fallback(typ), implicit=typ.implicit) + return typ + + return _type_asdict_inner(typ, seen_dataclasses=frozenset()) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 55a9a469e97b4..3efc861787f63 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -23,6 +23,7 @@ class DefaultPlugin(Plugin): def get_function_hook(self, fullname: str ) -> Optional[Callable[[FunctionContext], Type]]: from mypy.plugins import ctypes + from mypy.plugins import dataclasses if fullname == 'contextlib.contextmanager': return contextmanager_callback @@ -30,6 +31,8 @@ def get_function_hook(self, fullname: str return open_callback elif fullname == 'ctypes.Array': return ctypes.array_constructor_callback + elif fullname == 'dataclasses.asdict': + return dataclasses.asdict_callback return None def get_method_signature_hook(self, fullname: str diff --git a/mypy/typeshed b/mypy/typeshed index 375e8c1127288..b71ccbb136671 160000 --- a/mypy/typeshed +++ b/mypy/typeshed @@ -1 +1 @@ -Subproject commit 375e8c11272889cbb653a9cf680ff9b889309105 +Subproject commit b71ccbb136671307c1d94bdda00b4543f67926d2 diff --git a/test-data/unit/check-dataclasses.test b/test-data/unit/check-dataclasses.test index 97bb9954ec293..6ce8e9498620f 100644 --- a/test-data/unit/check-dataclasses.test +++ b/test-data/unit/check-dataclasses.test @@ -1032,3 +1032,292 @@ class B(A): reveal_type(B) # N: Revealed type is 'def (foo: builtins.int) -> __main__.B' [builtins fixtures/property.pyi] + +[case testDataclassesAsdict] +from dataclasses import dataclass, asdict + +@dataclass +class Person: + name: str + age: int + +@dataclass +class NotQuiteAPerson: + name: str + other_field: str + + +reveal_type(Person) # N: Revealed type is 'def (name: builtins.str, age: builtins.int) -> __main__.Person' +reveal_type(asdict(Person('John', 32))) # N: Revealed type is 'TypedDict({'name': builtins.str, 'age': builtins.int})' +Person(**asdict(Person('John', 32))) # Round-trip works +NotQuiteAPerson(**asdict(Person('John', 32))) # E: Extra argument "age" from **args for "NotQuiteAPerson" + +reveal_type(asdict(Person('John', 32))) # N: Revealed type is 'TypedDict({'name': builtins.str, 'age': builtins.int})' + +[typing fixtures/typing-full.pyi] + +[case testDataclassesAsdict] +from dataclasses import dataclass, asdict + +# TODO: Recursion, Unions, subclasses + +@dataclass +class Person: + name: str + age: int + +reveal_type(asdict(Person('John', 32))) # N: Revealed type is 'TypedDict({'name': builtins.str, 'age': builtins.int})' + +[typing fixtures/typing-full.pyi] + +[case testDataclassesAsdictRecursion] +from dataclasses import dataclass, asdict +from typing import Optional + +@dataclass +class C: + a: 'A' + +@dataclass +class B: + c: C + +@dataclass +class A: + b: Optional[B] = None + +# Recursion is not supported, so fall back +result = asdict(A(B(C(A())))) # E: Recursive types are not supported in call to asdict, so falling back to Dict[str, Any] +reveal_type(result) # N: Revealed type is 'TypedDict({'b': Union[TypedDict({'c': TypedDict({'a': builtins.dict[builtins.str, Any]})}), None]})' + +[typing fixtures/typing-full.pyi] +[builtins fixtures/tuple.pyi] +[builtins fixtures/dict.pyi] + +[case testDataclassesAsdictUnions] +from dataclasses import dataclass, asdict +from typing import Union + +@dataclass +class Card: + last4: int + +@dataclass +class Customer: + card: Union[str, Card] + +reveal_type(asdict(Customer("foo"))) # N: Revealed type is 'TypedDict({'card': Union[builtins.str, TypedDict({'last4': builtins.int})]})' + +[typing fixtures/typing-full.pyi] + +[case testDataclassesAsdictDictFactory] +from dataclasses import dataclass, asdict +from typing import List, Tuple, Dict, Any + +@dataclass +class Person: + name: str + age: int + +def my_dict_factory(seq: List[Tuple[str, Any]]) -> Dict[str, Any]: + pass + +# Passing in any a dict_factory falls back to default return value +reveal_type(asdict(Person('John', 32), dict_factory=my_dict_factory)) # N: Revealed type is 'builtins.dict*[builtins.str, Any]' + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dict.pyi] + +[case testDataclassesAsdictList] +from dataclasses import dataclass, asdict +from typing import List, Any + +@dataclass +class Person: + name: str + age: int + + +@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=[], +) +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]' + +[typing fixtures/typing-full.pyi] +[builtins fixtures/list.pyi] + + +[case testDataclassesAsdictListSubclass] +from dataclasses import dataclass, asdict +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: + list_subclass: MyList[Person] + list_subclass_2_typevars: MyListWith2TypeVars[Person, int] + list_subclass_with_constraint: MyListWithConstraint[Person] + +instance = Course( + list_subclass=MyList([]), + list_subclass_2_typevars=MyListWith2TypeVars[Person, int]([Person("John", 23)]), + list_subclass_with_constraint=MyListWithConstraint([Person("Tim", 29)]) +) +result = asdict(instance) + +# Supertypes (list) are returned, since there could be a constraint on the TypeVar +# used on the subclass such that when the type argument to the subclass is substituted with a TypedDict, +# it may not type-check. +reveal_type(result['list_subclass']) # N: Revealed type is 'builtins.list[TypedDict({'name': builtins.str, 'age': builtins.int})]' +reveal_type(result['list_subclass_2_typevars']) # N: Revealed type is 'builtins.list[TypedDict({'name': builtins.str, 'age': builtins.int})]' +reveal_type(result['list_subclass_with_constraint']) # N: Revealed type is 'builtins.list[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 + +@dataclass +class Person: + name: str + age: int + +@dataclass +class Course: + participants_by_name: Dict[str, Person] + +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})]' + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dict.pyi] + +[case testDataclassesAsdictDictSubclass] +from dataclasses import dataclass, asdict +from typing import Dict, Generic, TypeVar + +@dataclass +class Person: + name: str + age: int + +_KT = TypeVar("_KT") +_VT = TypeVar("_VT") +_Other = TypeVar("_Other") +class MyDict(Dict[_KT, _VT], Generic[_Other, _KT, _VT]): + pass + +@dataclass +class Course: + participants_by_name: MyDict[int, str, Person] + +instance = Course(participants_by_name=MyDict[int, str, Person]([("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})]' + +[typing fixtures/typing-full.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 + + def staff_method(self): + pass + +@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): + # It's unclear what the behavior should be here. + # At runtime, this isinstance check succeeds, but it is not sound, because + # the instance of Staff would have any field that is a dataclass replaced with dicts (TypedDicts), + # and there could be other transformations (e.g. to collection types) such that the resultant type is + # not compatible with the original. + # So, to be conservative, we are saying that staff is a "new" type with a fallback of type Staff. + # This reveal_type should produce no output by the type-checker since the code is unreachable, due to + # staff NOT being an instance of Staff. + reveal_type(staff) + +reveal_type(staff[0]) # N: Revealed type is 'TypedDict({'name': builtins.str, 'age': builtins.int})' + +# TODO: Make this work: +# This is OK, because although staff is not an instance of Staff, +# it has Staff as a fallback, so method calls should work +# staff.staff_method() +# TODO: Make this work: +# reveal_type(staff.teacher) # Revealed type is 'TypedDict({'name': builtins.str, 'age': builtins.int})' + +[typing fixtures/typing-full.pyi] +[builtins fixtures/tuple.pyi] diff --git a/test-data/unit/lib-stub/dataclasses.pyi b/test-data/unit/lib-stub/dataclasses.pyi index 160cfcd066ba9..526a9831552de 100644 --- a/test-data/unit/lib-stub/dataclasses.pyi +++ b/test-data/unit/lib-stub/dataclasses.pyi @@ -1,4 +1,4 @@ -from typing import Any, Callable, Generic, Mapping, Optional, TypeVar, overload, Type +from typing import Any, Callable, Generic, Mapping, Optional, TypeVar, overload, Type, Dict, List, Tuple _T = TypeVar('_T') @@ -6,6 +6,11 @@ class InitVar(Generic[_T]): ... +@overload +def asdict(obj: Any) -> Dict[str, Any]: ... +@overload +def asdict(obj: Any, *, dict_factory: Callable[[List[Tuple[str, Any]]], _T]) -> _T: ... + @overload def dataclass(_cls: Type[_T]) -> Type[_T]: ...