From 7beaec2e4a1c7891b044b45e538a472dbe86f240 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Wed, 5 Apr 2023 13:04:21 +0300 Subject: [PATCH] Support descriptors in dataclass transform (#15006) Infer `__init__` argument types from the signatures of descriptor `__set__` methods, if present. We can't (easily) perform full type inference in a plugin, so we cheat and use a simplified implementation that should still cover most use cases. Here we assume that `__set__` is not decorated or overloaded, in particular. Fixes #14868. --- mypy/plugins/dataclasses.py | 61 ++++- test-data/unit/check-dataclass-transform.test | 214 ++++++++++++++++++ 2 files changed, 269 insertions(+), 6 deletions(-) diff --git a/mypy/plugins/dataclasses.py b/mypy/plugins/dataclasses.py index 16b1595e3cb8..e84e1dbb9491 100644 --- a/mypy/plugins/dataclasses.py +++ b/mypy/plugins/dataclasses.py @@ -6,7 +6,7 @@ from typing_extensions import Final from mypy import errorcodes, message_registry -from mypy.expandtype import expand_type +from mypy.expandtype import expand_type, expand_type_by_instance from mypy.nodes import ( ARG_NAMED, ARG_NAMED_OPT, @@ -23,6 +23,7 @@ Context, DataclassTransformSpec, Expression, + FuncDef, IfStmt, JsonDict, NameExpr, @@ -98,7 +99,7 @@ def __init__( self.has_default = has_default self.line = line self.column = column - self.type = type + self.type = type # Type as __init__ argument self.info = info self.kw_only = kw_only self.is_neither_frozen_nor_nonfrozen = is_neither_frozen_nor_nonfrozen @@ -535,9 +536,12 @@ def collect_attributes(self) -> list[DataclassAttribute] | None: elif not isinstance(stmt.rvalue, TempNode): has_default = True - if not has_default: - # Make all non-default attributes implicit because they are de-facto set - # on self in the generated __init__(), not in the class body. + if not has_default and self._spec is _TRANSFORM_SPEC_FOR_DATACLASSES: + # Make all non-default dataclass attributes implicit because they are de-facto + # set on self in the generated __init__(), not in the class body. On the other + # hand, we don't know how custom dataclass transforms initialize attributes, + # so we don't treat them as implicit. This is required to support descriptors + # (https://github.com/python/mypy/issues/14868). sym.implicit = True is_kw_only = kw_only @@ -578,6 +582,7 @@ def collect_attributes(self) -> list[DataclassAttribute] | None: ) current_attr_names.add(lhs.name) + init_type = self._infer_dataclass_attr_init_type(sym, lhs.name, stmt) found_attrs[lhs.name] = DataclassAttribute( name=lhs.name, alias=alias, @@ -586,7 +591,7 @@ def collect_attributes(self) -> list[DataclassAttribute] | None: has_default=has_default, line=stmt.line, column=stmt.column, - type=sym.type, + type=init_type, info=cls.info, kw_only=is_kw_only, is_neither_frozen_nor_nonfrozen=_has_direct_dataclass_transform_metaclass( @@ -755,6 +760,50 @@ def _get_bool_arg(self, name: str, default: bool) -> bool: return require_bool_literal_argument(self._api, expression, name, default) return default + def _infer_dataclass_attr_init_type( + self, sym: SymbolTableNode, name: str, context: Context + ) -> Type | None: + """Infer __init__ argument type for an attribute. + + In particular, possibly use the signature of __set__. + """ + default = sym.type + if sym.implicit: + return default + t = get_proper_type(sym.type) + + # Perform a simple-minded inference from the signature of __set__, if present. + # We can't use mypy.checkmember here, since this plugin runs before type checking. + # We only support some basic scanerios here, which is hopefully sufficient for + # the vast majority of use cases. + if not isinstance(t, Instance): + return default + setter = t.type.get("__set__") + if setter: + if isinstance(setter.node, FuncDef): + super_info = t.type.get_containing_type_info("__set__") + assert super_info + if setter.type: + setter_type = get_proper_type( + map_type_from_supertype(setter.type, t.type, super_info) + ) + else: + return AnyType(TypeOfAny.unannotated) + if isinstance(setter_type, CallableType) and setter_type.arg_kinds == [ + ARG_POS, + ARG_POS, + ARG_POS, + ]: + return expand_type_by_instance(setter_type.arg_types[2], t) + else: + self._api.fail( + f'Unsupported signature for "__set__" in "{t.type.name}"', context + ) + else: + self._api.fail(f'Unsupported "__set__" in "{t.type.name}"', context) + + return default + def add_dataclass_tag(info: TypeInfo) -> None: # The value is ignored, only the existence matters. diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test index 8d8e38997582..0b696d36b9bd 100644 --- a/test-data/unit/check-dataclass-transform.test +++ b/test-data/unit/check-dataclass-transform.test @@ -807,3 +807,217 @@ reveal_type(bar.base) # N: Revealed type is "builtins.int" [typing fixtures/typing-full.pyi] [builtins fixtures/dataclasses.pyi] + +[case testDataclassTransformSimpleDescriptor] +# flags: --python-version 3.11 + +from typing import dataclass_transform, overload, Any + +@dataclass_transform() +def my_dataclass(cls): ... + +class Desc: + @overload + def __get__(self, instance: None, owner: Any) -> Desc: ... + @overload + def __get__(self, instance: object, owner: Any) -> str: ... + def __get__(self, instance: object | None, owner: Any) -> Desc | str: ... + + def __set__(self, instance: Any, value: str) -> None: ... + +@my_dataclass +class C: + x: Desc + y: int + +C(x='x', y=1) +C(x=1, y=1) # E: Argument "x" to "C" has incompatible type "int"; expected "str" +reveal_type(C(x='x', y=1).x) # N: Revealed type is "builtins.str" +reveal_type(C(x='x', y=1).y) # N: Revealed type is "builtins.int" +reveal_type(C.x) # N: Revealed type is "__main__.Desc" + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] + +[case testDataclassTransformUnannotatedDescriptor] +# flags: --python-version 3.11 + +from typing import dataclass_transform, overload, Any + +@dataclass_transform() +def my_dataclass(cls): ... + +class Desc: + @overload + def __get__(self, instance: None, owner: Any) -> Desc: ... + @overload + def __get__(self, instance: object, owner: Any) -> str: ... + def __get__(self, instance: object | None, owner: Any) -> Desc | str: ... + + def __set__(*args, **kwargs): ... + +@my_dataclass +class C: + x: Desc + y: int + +C(x='x', y=1) +C(x=1, y=1) +reveal_type(C(x='x', y=1).x) # N: Revealed type is "builtins.str" +reveal_type(C(x='x', y=1).y) # N: Revealed type is "builtins.int" +reveal_type(C.x) # N: Revealed type is "__main__.Desc" + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] + +[case testDataclassTransformGenericDescriptor] +# flags: --python-version 3.11 + +from typing import dataclass_transform, overload, Any, TypeVar, Generic + +@dataclass_transform() +def my_dataclass(frozen: bool = False): ... + +T = TypeVar("T") + +class Desc(Generic[T]): + @overload + def __get__(self, instance: None, owner: Any) -> Desc[T]: ... + @overload + def __get__(self, instance: object, owner: Any) -> T: ... + def __get__(self, instance: object | None, owner: Any) -> Desc | T: ... + + def __set__(self, instance: Any, value: T) -> None: ... + +@my_dataclass() +class C: + x: Desc[str] + +C(x='x') +C(x=1) # E: Argument "x" to "C" has incompatible type "int"; expected "str" +reveal_type(C(x='x').x) # N: Revealed type is "builtins.str" +reveal_type(C.x) # N: Revealed type is "__main__.Desc[builtins.str]" + +@my_dataclass() +class D(C): + y: Desc[int] + +d = D(x='x', y=1) +reveal_type(d.x) # N: Revealed type is "builtins.str" +reveal_type(d.y) # N: Revealed type is "builtins.int" +reveal_type(D.x) # N: Revealed type is "__main__.Desc[builtins.str]" +reveal_type(D.y) # N: Revealed type is "__main__.Desc[builtins.int]" + +@my_dataclass(frozen=True) +class F: + x: Desc[str] = Desc() + +F(x='x') +F(x=1) # E: Argument "x" to "F" has incompatible type "int"; expected "str" +reveal_type(F(x='x').x) # N: Revealed type is "builtins.str" +reveal_type(F.x) # N: Revealed type is "__main__.Desc[builtins.str]" + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] + +[case testDataclassTransformGenericDescriptorWithInheritance] +# flags: --python-version 3.11 + +from typing import dataclass_transform, overload, Any, TypeVar, Generic + +@dataclass_transform() +def my_dataclass(cls): ... + +T = TypeVar("T") + +class Desc(Generic[T]): + @overload + def __get__(self, instance: None, owner: Any) -> Desc[T]: ... + @overload + def __get__(self, instance: object, owner: Any) -> T: ... + def __get__(self, instance: object | None, owner: Any) -> Desc | T: ... + + def __set__(self, instance: Any, value: T) -> None: ... + +class Desc2(Desc[str]): + pass + +@my_dataclass +class C: + x: Desc2 + +C(x='x') +C(x=1) # E: Argument "x" to "C" has incompatible type "int"; expected "str" +reveal_type(C(x='x').x) # N: Revealed type is "builtins.str" +reveal_type(C.x) # N: Revealed type is "__main__.Desc[builtins.str]" + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] + +[case testDataclassTransformDescriptorWithDifferentGetSetTypes] +# flags: --python-version 3.11 + +from typing import dataclass_transform, overload, Any + +@dataclass_transform() +def my_dataclass(cls): ... + +class Desc: + @overload + def __get__(self, instance: None, owner: Any) -> int: ... + @overload + def __get__(self, instance: object, owner: Any) -> str: ... + def __get__(self, instance, owner): ... + + def __set__(self, instance: Any, value: bytes) -> None: ... + +@my_dataclass +class C: + x: Desc + +c = C(x=b'x') +C(x=1) # E: Argument "x" to "C" has incompatible type "int"; expected "bytes" +reveal_type(c.x) # N: Revealed type is "builtins.str" +reveal_type(C.x) # N: Revealed type is "builtins.int" +c.x = b'x' +c.x = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "bytes") + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] + +[case testDataclassTransformUnsupportedDescriptors] +# flags: --python-version 3.11 + +from typing import dataclass_transform, overload, Any + +@dataclass_transform() +def my_dataclass(cls): ... + +class Desc: + @overload + def __get__(self, instance: None, owner: Any) -> int: ... + @overload + def __get__(self, instance: object, owner: Any) -> str: ... + def __get__(self, instance, owner): ... + + def __set__(*args, **kwargs) -> None: ... + +class Desc2: + @overload + def __get__(self, instance: None, owner: Any) -> int: ... + @overload + def __get__(self, instance: object, owner: Any) -> str: ... + def __get__(self, instance, owner): ... + + @overload + def __set__(self, instance: Any, value: bytes) -> None: ... + @overload + def __set__(self) -> None: ... + def __set__(self, *args, **kawrga) -> None: ... + +@my_dataclass +class C: + x: Desc # E: Unsupported signature for "__set__" in "Desc" + y: Desc2 # E: Unsupported "__set__" in "Desc2" +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi]