[dataclass_transform] support function overloads (#14651)

Extends the existing decorator support to include overloads. This
doesn't require much extra work: we just update
`find_dataclass_transform_spec` to search for the first overload
decorated with `dataclass_transform()`.

mypy/semanal_shared.py

@@ -19,6 +19,7 @@
     Expression,
     FuncDef,
     Node,
+    OverloadedFuncDef,
     RefExpr,
     SymbolNode,
     SymbolTable,
@@ -379,6 +380,17 @@ def find_dataclass_transform_spec(node: Node | None) -> DataclassTransformSpec |
         # `@dataclass_transform(...)` syntax and never `@dataclass_transform`
         node = node.func
 
+    if isinstance(node, OverloadedFuncDef):
+        # The dataclass_transform decorator may be attached to any single overload, so we must
+        # search them all.
+        # Note that using more than one decorator is undefined behavior, so we can just take the
+        # first that we find.
+        for candidate in node.items:
+            spec = find_dataclass_transform_spec(candidate)
+            if spec is not None:
+                return spec
+        return find_dataclass_transform_spec(node.impl)
+
     # For functions, we can directly consult the AST field for the spec
     if isinstance(node, FuncDef):
         return node.dataclass_transform_spec

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

@@ -210,6 +210,60 @@ Foo(5)
 [typing fixtures/typing-full.pyi]
 [builtins fixtures/dataclasses.pyi]
 
+[case testDataclassTransformOverloadsDecoratorOnOverload]
+# flags: --python-version 3.11
+from typing import dataclass_transform, overload, Any, Callable, Type, Literal
+
+@overload
+def my_dataclass(*, foo: str) -> Callable[[Type], Type]: ...
+@overload
+@dataclass_transform(frozen_default=True)
+def my_dataclass(*, foo: int) -> Callable[[Type], Type]: ...
+def my_dataclass(*, foo: Any) -> Callable[[Type], Type]:
+    return lambda cls: cls
+@my_dataclass(foo="hello")
+class A:
+    a: int
+@my_dataclass(foo=5)
+class B:
+    b: int
+
+reveal_type(A)  # N: Revealed type is "def (a: builtins.int) -> __main__.A"
+reveal_type(B)  # N: Revealed type is "def (b: builtins.int) -> __main__.B"
+A(1, "hello")  # E: Too many arguments for "A"
+a = A(1)
+a.a = 2  # E: Property "a" defined in "A" is read-only
+
+[typing fixtures/typing-full.pyi]
+[builtins fixtures/dataclasses.pyi]
+
+[case testDataclassTransformOverloadsDecoratorOnImpl]
+# flags: --python-version 3.11
+from typing import dataclass_transform, overload, Any, Callable, Type, Literal
+
+@overload
+def my_dataclass(*, foo: str) -> Callable[[Type], Type]: ...
+@overload
+def my_dataclass(*, foo: int) -> Callable[[Type], Type]: ...
+@dataclass_transform(frozen_default=True)
+def my_dataclass(*, foo: Any) -> Callable[[Type], Type]:
+    return lambda cls: cls
+@my_dataclass(foo="hello")
+class A:
+    a: int
+@my_dataclass(foo=5)
+class B:
+    b: int
+
+reveal_type(A)  # N: Revealed type is "def (a: builtins.int) -> __main__.A"
+reveal_type(B)  # N: Revealed type is "def (b: builtins.int) -> __main__.B"
+A(1, "hello")  # E: Too many arguments for "A"
+a = A(1)
+a.a = 2  # E: Property "a" defined in "A" is read-only
+
+[typing fixtures/typing-full.pyi]
+[builtins fixtures/dataclasses.pyi]
+
 [case testDataclassTransformViaBaseClass]
 # flags: --python-version 3.11
 from typing import dataclass_transform