From 0d87c8dc6e436d268a9dcd0bede7e0f44eb7e43e Mon Sep 17 00:00:00 2001 From: Wesley Wright Date: Fri, 20 Jan 2023 19:33:33 +0000 Subject: [PATCH 1/5] [dataclass_transform] minimal integration of dataclass_transform with dataclasses plugin This is a very simple first step to implementing [PEP 0681](https://peps.python.org/pep-0681/#decorator-function-example), which will allow MyPy to recognize user-defined types that behave similarly to dataclasses. This initial implementation is very limited: we only support decorator-style use of `typing.dataclass_transform` and do not support passing additional options to the transform (such as `freeze` or `init`). Within MyPy, we add a new `is_dataclass_transform` field to `FuncBase` which is populated during semantic analysis. When we check for plugin hooks later, we add new special cases to use the dataclasses plugin if a class decorator is marked with `is_dataclass_transform`. Ideally we would use a proper plugin API; the hacky special case here can be replaced in subsequent iterations. --- mypy/nodes.py | 10 ++++++++- mypy/plugins/common.py | 2 +- mypy/semanal.py | 22 +++++++++++++++---- mypy/semanal_main.py | 15 +++++++++++-- mypy/semanal_shared.py | 5 +++++ test-data/unit/check-dataclass-transform.test | 22 +++++++++++++++++++ test-data/unit/fixtures/typing-medium.pyi | 2 ++ 7 files changed, 70 insertions(+), 8 deletions(-) create mode 100644 test-data/unit/check-dataclass-transform.test diff --git a/mypy/nodes.py b/mypy/nodes.py index 38639d553b3d..245367dc3044 100644 --- a/mypy/nodes.py +++ b/mypy/nodes.py @@ -480,7 +480,13 @@ def accept(self, visitor: StatementVisitor[T]) -> T: return visitor.visit_import_all(self) -FUNCBASE_FLAGS: Final = ["is_property", "is_class", "is_static", "is_final"] +FUNCBASE_FLAGS: Final = [ + "is_property", + "is_class", + "is_static", + "is_final", + "is_dataclass_transform", +] class FuncBase(Node): @@ -506,6 +512,7 @@ class FuncBase(Node): "is_static", # Uses "@staticmethod" "is_final", # Uses "@final" "_fullname", + "is_dataclass_transform", ) def __init__(self) -> None: @@ -524,6 +531,7 @@ def __init__(self) -> None: self.is_final = False # Name with module prefix self._fullname = "" + self.is_dataclass_transform = False @property @abstractmethod diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py index 07cd5dc7de7f..a2a38f256da3 100644 --- a/mypy/plugins/common.py +++ b/mypy/plugins/common.py @@ -19,7 +19,7 @@ Var, ) from mypy.plugin import CheckerPluginInterface, ClassDefContext, SemanticAnalyzerPluginInterface -from mypy.semanal import ALLOW_INCOMPATIBLE_OVERRIDE, set_callable_name +from mypy.semanal_shared import ALLOW_INCOMPATIBLE_OVERRIDE, set_callable_name from mypy.typeops import ( # noqa: F401 # Part of public API try_getting_str_literals as try_getting_str_literals, ) diff --git a/mypy/semanal.py b/mypy/semanal.py index 34cb45194d19..2b8152ccbd5a 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -194,6 +194,7 @@ Plugin, SemanticAnalyzerPluginInterface, ) +from mypy.plugins import dataclasses as dataclasses_plugin from mypy.reachability import ( ALWAYS_FALSE, ALWAYS_TRUE, @@ -208,6 +209,7 @@ from mypy.semanal_namedtuple import NamedTupleAnalyzer from mypy.semanal_newtype import NewTypeAnalyzer from mypy.semanal_shared import ( + ALLOW_INCOMPATIBLE_OVERRIDE, PRIORITY_FALLBACKS, SemanticAnalyzerInterface, calculate_tuple_fallback, @@ -304,10 +306,6 @@ # available very early on. CORE_BUILTIN_CLASSES: Final = ["object", "bool", "function"] -# Subclasses can override these Var attributes with incompatible types. This can also be -# set for individual attributes using 'allow_incompatible_override' of Var. -ALLOW_INCOMPATIBLE_OVERRIDE: Final = ("__slots__", "__deletable__", "__match_args__") - # Used for tracking incomplete references Tag: _TypeAlias = int @@ -1505,6 +1503,10 @@ def visit_decorator(self, dec: Decorator) -> None: removed.append(i) else: self.fail("@final cannot be used with non-method functions", d) + elif isinstance(d, CallExpr) and refers_to_fullname( + d.callee, "typing.dataclass_transform" + ): + dec.func.is_dataclass_transform = True elif not dec.var.is_property: # We have seen a "non-trivial" decorator before seeing @property, if # we will see a @property later, give an error, as we don't support this. @@ -1706,6 +1708,11 @@ def apply_class_plugin_hooks(self, defn: ClassDef) -> None: decorator_name = self.get_fullname_for_hook(decorator) if decorator_name: hook = self.plugin.get_class_decorator_hook(decorator_name) + # Special case: if the decorator is itself decorated with + # typing.dataclass_transform, apply the hook for the dataclasses plugin + # TODO: remove special casing here + if hook is None and is_dataclass_transform_decorator(decorator): + hook = dataclasses_plugin.dataclass_tag_callback if hook: hook(ClassDefContext(defn, decorator, self)) @@ -6596,3 +6603,10 @@ def halt(self, reason: str = ...) -> NoReturn: return isinstance(stmt, PassStmt) or ( isinstance(stmt, ExpressionStmt) and isinstance(stmt.expr, EllipsisExpr) ) + + +def is_dataclass_transform_decorator(node: Node | None) -> bool: + if isinstance(node, RefExpr): + return is_dataclass_transform_decorator(node.node) + + return isinstance(node, Decorator) and node.func.is_dataclass_transform diff --git a/mypy/semanal_main.py b/mypy/semanal_main.py index 9e3aeaa7fa4b..01515018d0e1 100644 --- a/mypy/semanal_main.py +++ b/mypy/semanal_main.py @@ -37,9 +37,11 @@ from mypy.nodes import Decorator, FuncDef, MypyFile, OverloadedFuncDef, TypeInfo, Var from mypy.options import Options from mypy.plugin import ClassDefContext +from mypy.plugins import dataclasses as dataclasses_plugin from mypy.semanal import ( SemanticAnalyzer, apply_semantic_analyzer_patches, + is_dataclass_transform_decorator, remove_imported_names_from_symtable, ) from mypy.semanal_classprop import ( @@ -450,11 +452,20 @@ def apply_hooks_to_class( ok = True for decorator in defn.decorators: with self.file_context(file_node, options, info): + hook = None + + decorator.accept(self) decorator_name = self.get_fullname_for_hook(decorator) if decorator_name: hook = self.plugin.get_class_decorator_hook_2(decorator_name) - if hook: - ok = ok and hook(ClassDefContext(defn, decorator, self)) + # Special case: if the decorator is itself decorated with + # typing.dataclass_transform, apply the hook for the dataclasses plugin + # TODO: remove special casing here + if hook is None and is_dataclass_transform_decorator(decorator): + hook = dataclasses_plugin.dataclass_class_maker_callback + + if hook: + ok = ok and hook(ClassDefContext(defn, decorator, self)) return ok diff --git a/mypy/semanal_shared.py b/mypy/semanal_shared.py index f4bc173b52d5..11c4af314a3b 100644 --- a/mypy/semanal_shared.py +++ b/mypy/semanal_shared.py @@ -38,6 +38,11 @@ get_proper_type, ) +# Subclasses can override these Var attributes with incompatible types. This can also be +# set for individual attributes using 'allow_incompatible_override' of Var. +ALLOW_INCOMPATIBLE_OVERRIDE: Final = ("__slots__", "__deletable__", "__match_args__") + + # Priorities for ordering of patches within the "patch" phase of semantic analysis # (after the main pass): diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test new file mode 100644 index 000000000000..49d0b79d051f --- /dev/null +++ b/test-data/unit/check-dataclass-transform.test @@ -0,0 +1,22 @@ +[case testDataclassTransformReusesDataclassLogic] +# flags: --python-version 3.7 +from typing import dataclass_transform, Type + +@dataclass_transform() +def my_dataclass(cls: Type) -> Type: + return cls + +@my_dataclass +class Person: + name: str + age: int + + def summary(self): + return "%s is %d years old." % (self.name, self.age) + +reveal_type(Person) # N: Revealed type is "def (name: builtins.str, age: builtins.int) -> __main__.Person" +Person('John', 32) +Person('Jonh', 21, None) # E: Too many arguments for "Person" + +[typing fixtures/typing-medium.pyi] +[builtins fixtures/dataclasses.pyi] diff --git a/test-data/unit/fixtures/typing-medium.pyi b/test-data/unit/fixtures/typing-medium.pyi index 863b0703989d..0d0e13468013 100644 --- a/test-data/unit/fixtures/typing-medium.pyi +++ b/test-data/unit/fixtures/typing-medium.pyi @@ -71,3 +71,5 @@ class ContextManager(Generic[T]): class _SpecialForm: pass TYPE_CHECKING = 1 + +def dataclass_transform() -> Callable[[T], T]: ... From 1c76978d5afd08f751b813567a7b936852747986 Mon Sep 17 00:00:00 2001 From: Wesley Collin Wright Date: Wed, 25 Jan 2023 19:38:32 +0000 Subject: [PATCH 2/5] also support typing_extensions.dataclass_transform --- mypy/semanal.py | 3 ++- mypy/types.py | 5 ++++ test-data/unit/check-dataclass-transform.test | 24 +++++++++++++++++++ test-data/unit/fixtures/dataclasses.pyi | 6 ++++- test-data/unit/lib-stub/typing_extensions.pyi | 2 ++ 5 files changed, 38 insertions(+), 2 deletions(-) diff --git a/mypy/semanal.py b/mypy/semanal.py index 2b8152ccbd5a..9f3b50eb21c9 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -236,6 +236,7 @@ from mypy.typeops import function_type, get_type_vars from mypy.types import ( ASSERT_TYPE_NAMES, + DATACLASS_TRANSFORM_NAMES, FINAL_DECORATOR_NAMES, FINAL_TYPE_NAMES, NEVER_NAMES, @@ -1504,7 +1505,7 @@ def visit_decorator(self, dec: Decorator) -> None: else: self.fail("@final cannot be used with non-method functions", d) elif isinstance(d, CallExpr) and refers_to_fullname( - d.callee, "typing.dataclass_transform" + d.callee, DATACLASS_TRANSFORM_NAMES ): dec.func.is_dataclass_transform = True elif not dec.var.is_property: diff --git a/mypy/types.py b/mypy/types.py index bf610a01b63b..74656cc270f3 100644 --- a/mypy/types.py +++ b/mypy/types.py @@ -150,6 +150,11 @@ "typing_extensions.Never", ) +DATACLASS_TRANSFORM_NAMES: Final = ( + "typing.dataclass_transform", + "typing_extensions.dataclass_transform", +) + # A placeholder used for Bogus[...] parameters _dummy: Final[Any] = object() diff --git a/test-data/unit/check-dataclass-transform.test b/test-data/unit/check-dataclass-transform.test index 49d0b79d051f..4f907e3186b6 100644 --- a/test-data/unit/check-dataclass-transform.test +++ b/test-data/unit/check-dataclass-transform.test @@ -20,3 +20,27 @@ Person('Jonh', 21, None) # E: Too many arguments for "Person" [typing fixtures/typing-medium.pyi] [builtins fixtures/dataclasses.pyi] + +[case testDataclassTransformIsFoundInTypingExtensions] +# flags: --python-version 3.7 +from typing import Type +from typing_extensions import dataclass_transform + +@dataclass_transform() +def my_dataclass(cls: Type) -> Type: + return cls + +@my_dataclass +class Person: + name: str + age: int + + def summary(self): + return "%s is %d years old." % (self.name, self.age) + +reveal_type(Person) # N: Revealed type is "def (name: builtins.str, age: builtins.int) -> __main__.Person" +Person('John', 32) +Person('Jonh', 21, None) # E: Too many arguments for "Person" + +[typing fixtures/typing-full.pyi] +[builtins fixtures/dataclasses.pyi] diff --git a/test-data/unit/fixtures/dataclasses.pyi b/test-data/unit/fixtures/dataclasses.pyi index 206843a88b24..7de40af9cfe7 100644 --- a/test-data/unit/fixtures/dataclasses.pyi +++ b/test-data/unit/fixtures/dataclasses.pyi @@ -37,7 +37,11 @@ class dict(Mapping[KT, VT]): def get(self, k: KT, default: Union[KT, _T]) -> Union[VT, _T]: pass def __len__(self) -> int: ... -class list(Generic[_T], Sequence[_T]): pass +class list(Generic[_T], Sequence[_T]): + def __contains__(self, item: object) -> int: pass + def __getitem__(self, key: int) -> _T: pass + def __iter__(self) -> Iterator[_T]: pass + class function: pass class classmethod: pass property = object() diff --git a/test-data/unit/lib-stub/typing_extensions.pyi b/test-data/unit/lib-stub/typing_extensions.pyi index cbf692fc7111..89f7108fe83c 100644 --- a/test-data/unit/lib-stub/typing_extensions.pyi +++ b/test-data/unit/lib-stub/typing_extensions.pyi @@ -57,3 +57,5 @@ class _TypedDict(Mapping[str, object]): def TypedDict(typename: str, fields: Dict[str, Type[_T]], *, total: Any = ...) -> Type[dict]: ... def reveal_type(__obj: T) -> T: pass + +def dataclass_transform() -> Callable[[T], T]: ... From 6a8e690b9439980e2de24cd50b1f89ddb95f26ee Mon Sep 17 00:00:00 2001 From: Wesley Collin Wright Date: Wed, 25 Jan 2023 22:27:57 +0000 Subject: [PATCH 3/5] document is_dataclass_transform --- mypy/nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypy/nodes.py b/mypy/nodes.py index 245367dc3044..c856841cd82d 100644 --- a/mypy/nodes.py +++ b/mypy/nodes.py @@ -512,7 +512,7 @@ class FuncBase(Node): "is_static", # Uses "@staticmethod" "is_final", # Uses "@final" "_fullname", - "is_dataclass_transform", + "is_dataclass_transform", # Is decorated with "@typing.dataclass_transform" or similar ) def __init__(self) -> None: From 8207a9ffaf73d1e8a7f164b09dfac4d61625b687 Mon Sep 17 00:00:00 2001 From: Wesley Collin Wright Date: Wed, 25 Jan 2023 22:28:27 +0000 Subject: [PATCH 4/5] no need to accept decorators again --- mypy/semanal_main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mypy/semanal_main.py b/mypy/semanal_main.py index 01515018d0e1..ab38566487b8 100644 --- a/mypy/semanal_main.py +++ b/mypy/semanal_main.py @@ -454,7 +454,6 @@ def apply_hooks_to_class( with self.file_context(file_node, options, info): hook = None - decorator.accept(self) decorator_name = self.get_fullname_for_hook(decorator) if decorator_name: hook = self.plugin.get_class_decorator_hook_2(decorator_name) From 824093186b6bd63f420074c8d502e6a57688501e Mon Sep 17 00:00:00 2001 From: Wesley Collin Wright Date: Thu, 26 Jan 2023 18:29:27 +0000 Subject: [PATCH 5/5] placate flake8 --- mypy/nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypy/nodes.py b/mypy/nodes.py index c856841cd82d..98976f4fe56a 100644 --- a/mypy/nodes.py +++ b/mypy/nodes.py @@ -512,7 +512,7 @@ class FuncBase(Node): "is_static", # Uses "@staticmethod" "is_final", # Uses "@final" "_fullname", - "is_dataclass_transform", # Is decorated with "@typing.dataclass_transform" or similar + "is_dataclass_transform", # Is decorated with "@typing.dataclass_transform" or similar ) def __init__(self) -> None: