[dataclass_transform] support class decorator parameters (#14561)

The initial implementation of `typing.dataclass_transform` only
supported the no-argument `@decorator` form; this adds support for the
`@decorator(...)` form supporting the same arguments we support for
`dataclasses.dataclass`. This also matches the list of arguments
specified in PEP 681.

Co-authored-by: Wesley Collin Wright <wesleyw@dropbox.com>

Author: wesleywright

mypy/plugins/common.py

@@ -13,6 +13,7 @@
     Expression,
     FuncDef,
     JsonDict,
+    Node,
     PassStmt,
     RefExpr,
     SymbolTableNode,
@@ -68,19 +69,7 @@ def _get_argument(call: CallExpr, name: str) -> Expression | None:
     #
     # Note: I'm not hard-coding the index so that in the future we can support other
     # attrib and class makers.
-    if not isinstance(call.callee, RefExpr):
-        return None
-
-    callee_type = None
-    callee_node = call.callee.node
-    if isinstance(callee_node, (Var, SYMBOL_FUNCBASE_TYPES)) and callee_node.type:
-        callee_node_type = get_proper_type(callee_node.type)
-        if isinstance(callee_node_type, Overloaded):
-            # We take the last overload.
-            callee_type = callee_node_type.items[-1]
-        elif isinstance(callee_node_type, CallableType):
-            callee_type = callee_node_type
-
+    callee_type = _get_callee_type(call)
     if not callee_type:
         return None
 
@@ -94,6 +83,31 @@ def _get_argument(call: CallExpr, name: str) -> Expression | None:
             return attr_value
         if attr_name == argument.name:
             return attr_value
+
+    return None
+
+
+def _get_callee_type(call: CallExpr) -> CallableType | None:
+    """Return the type of the callee, regardless of its syntatic form."""
+
+    callee_node: Node | None = call.callee
+
+    if isinstance(callee_node, RefExpr):
+        callee_node = callee_node.node
+
+    # Some decorators may be using typing.dataclass_transform, which is itself a decorator, so we
+    # need to unwrap them to get at the true callee
+    if isinstance(callee_node, Decorator):
+        callee_node = callee_node.func
+
+    if isinstance(callee_node, (Var, SYMBOL_FUNCBASE_TYPES)) and callee_node.type:
+        callee_node_type = get_proper_type(callee_node.type)
+        if isinstance(callee_node_type, Overloaded):
+            # We take the last overload.
+            return callee_node_type.items[-1]
+        elif isinstance(callee_node_type, CallableType):
+            return callee_node_type
+
     return None
 
 

mypy/semanal.py

@@ -6646,5 +6646,16 @@ def halt(self, reason: str = ...) -> NoReturn:
 def is_dataclass_transform_decorator(node: Node | None) -> bool:
     if isinstance(node, RefExpr):
         return is_dataclass_transform_decorator(node.node)
+    if isinstance(node, CallExpr):
+        # Like dataclasses.dataclass, transform-based decorators can be applied either with or
+        # without parameters; ie, both of these forms are accepted:
+        #
+        # @typing.dataclass_transform
+        # class Foo: ...
+        # @typing.dataclass_transform(eq=True, order=True, ...)
+        # class Bar: ...
+        #
+        # We need to unwrap the call for the second variant.
+        return is_dataclass_transform_decorator(node.callee)
 
     return isinstance(node, Decorator) and node.func.is_dataclass_transform

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

@@ -44,3 +44,44 @@ Person('Jonh', 21, None)  # E: Too many arguments for "Person"
 
 [typing fixtures/typing-full.pyi]
 [builtins fixtures/dataclasses.pyi]
+
+[case testDataclassTransformParametersAreApplied]
+# flags: --python-version 3.7
+from typing import dataclass_transform, Callable, Type
+
+@dataclass_transform()
+def my_dataclass(*, eq: bool, order: bool) -> Callable[[Type], Type]:
+    def transform(cls: Type) -> Type:
+        return cls
+    return transform
+
+@my_dataclass(eq=False, order=True)
+class Person:  # E: eq must be True if order is True
+    name: str
+    age: int
+
+reveal_type(Person)  # N: Revealed type is "def (name: builtins.str, age: builtins.int) -> __main__.Person"
+Person('John', 32)
+Person('John', 21, None)  # E: Too many arguments for "Person"
+
+[typing fixtures/typing-medium.pyi]
+[builtins fixtures/dataclasses.pyi]
+
+[case testDataclassTransformParametersMustBeBoolLiterals]
+# flags: --python-version 3.7
+from typing import dataclass_transform, Callable, Type
+
+@dataclass_transform()
+def my_dataclass(*, eq: bool = True, order: bool = False) -> Callable[[Type], Type]:
+    def transform(cls: Type) -> Type:
+        return cls
+    return transform
+
+BOOL_CONSTANT = True
+@my_dataclass(eq=BOOL_CONSTANT)  # E: "eq" argument must be True or False.
+class A: ...
+@my_dataclass(order=not False)  # E: "order" argument must be True or False.
+class B: ...
+
+[typing fixtures/typing-medium.pyi]
+[builtins fixtures/dataclasses.pyi]