[dataclass_transform] minimal implementation of dataclass_transform (#14523)

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 existing 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.

Co-authored-by: Wesley Wright <[email protected]>

Author: wesleywright

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",  # Is decorated with "@typing.dataclass_transform" or similar
     )
 
     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

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,
 )

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,
@@ -234,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,
@@ -304,10 +307,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
@@ -1508,6 +1507,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, DATACLASS_TRANSFORM_NAMES
+            ):
+                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.
@@ -1709,6 +1712,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))
 
@@ -6599,3 +6607,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

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 (
@@ -457,11 +459,19 @@ def apply_hooks_to_class(
     ok = True
     for decorator in defn.decorators:
         with self.file_context(file_node, options, info):
+            hook = None
+
             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
 
 

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):
 

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()
 

test-data/unit/check-dataclass-transform.test

@@ -0,0 +1,46 @@
+[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]
+
+[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]

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()

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]: ...

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]: ...