stubtest: attempt to resolve decorators from their type (#20867)

Fixes #19689

Author: hauntsaninja

mypy/stubtest.py

@@ -12,7 +12,6 @@
 import enum
 import functools
 import importlib
-import importlib.machinery
 import inspect
 import os
 import pkgutil
@@ -35,11 +34,11 @@
 
 import mypy.build
 import mypy.checkexpr
-import mypy.checkmember
 import mypy.erasetype
 import mypy.modulefinder
 import mypy.nodes
 import mypy.state
+import mypy.subtypes
 import mypy.types
 import mypy.version
 from mypy import nodes
@@ -1529,7 +1528,7 @@ def apply_decorator_to_funcitem(
         ):
             return func
         if decorator.fullname == "builtins.classmethod":
-            if func.arguments[0].variable.name not in ("cls", "mcs", "metacls"):
+            if func.arguments[0].variable.name not in ("_cls", "cls", "mcs", "metacls"):
                 raise StubtestFailure(
                     f"unexpected class parameter name {func.arguments[0].variable.name!r} "
                     f"in {dec.fullname}"
@@ -1547,11 +1546,70 @@ def apply_decorator_to_funcitem(
     for decorator in dec.original_decorators:
         resulting_func = apply_decorator_to_funcitem(decorator, func)
         if resulting_func is None:
+            # We couldn't figure out how to apply the decorator by transforming nodes, so try to
+            # reconstitute a FuncDef from the resulting type of the decorator
+            # This is worse because e.g. we lose the values of defaults
+            dec_type = mypy.types.get_proper_type(dec.type)
+            callable_type = None
+            if isinstance(dec_type, mypy.types.Instance):
+                callable_type = mypy.subtypes.find_member(
+                    "__call__", dec_type, dec_type, is_operator=True
+                )
+            elif isinstance(dec_type, mypy.types.CallableType):
+                callable_type = dec_type
+
+            callable_type = mypy.types.get_proper_type(callable_type)
+            if isinstance(callable_type, mypy.types.CallableType):
+                return _resolve_funcitem_from_callable_type(dec, callable_type)
             return None
+
         func = resulting_func
     return func
 
 
+def _resolve_funcitem_from_callable_type(
+    dec: nodes.Decorator, typ: mypy.types.CallableType
+) -> nodes.FuncDef | None:
+    if (
+        typ.arg_kinds == [nodes.ARG_STAR, nodes.ARG_STAR2]
+        and (var_arg := typ.var_arg()) is not None
+        and isinstance(mypy.types.get_proper_type(var_arg.typ), mypy.types.AnyType)
+        and (var_kwarg := typ.kw_arg()) is not None
+        and isinstance(mypy.types.get_proper_type(var_kwarg.typ), mypy.types.AnyType)
+    ):
+        # There isn't a FuncDef we can invent corresponding to a Callable[..., T]
+        return None
+
+    args: list[nodes.Argument] = []
+    for i, (arg_type, arg_kind, arg_name) in enumerate(
+        zip(typ.arg_types, typ.arg_kinds, typ.arg_names, strict=True)
+    ):
+        var_name = arg_name if arg_name is not None else f"__arg{i}"
+        var = nodes.Var(var_name, arg_type)
+        pos_only = arg_name is None and arg_kind == nodes.ARG_POS
+        args.append(
+            nodes.Argument(
+                variable=var,
+                type_annotation=arg_type,
+                initializer=None,  # CallableType doesn't store the values of defaults
+                kind=arg_kind,
+                pos_only=pos_only,
+            )
+        )
+
+    if dec.func.is_class:
+        if not args:
+            return None
+        # Munge classmethods, similar to logic in _resolve_funcitem_from_decorator
+        if args[0].variable.name not in ("_cls", "cls", "mcs", "metacls"):
+            return None
+        args.pop(0)
+
+    ret = nodes.FuncDef(name=typ.name or "", arguments=args, body=nodes.Block([]), typ=typ)
+    ret.is_class = dec.func.is_class
+    return ret
+
+
 @verify.register(nodes.Decorator)
 def verify_decorator(
     stub: nodes.Decorator, runtime: MaybeMissing[Any], object_path: list[str]

mypy/test/teststubtest.py

@@ -900,6 +900,56 @@ def f(a, *args): ...
             error=None,
         )
 
+    @collect_cases
+    def test_decorated_overload(self) -> Iterator[Case]:
+        yield Case(
+            stub="""
+            from typing import overload
+
+            class _dec1:
+                def __init__(self, func: object) -> None: ...
+                def __call__(self, x: str) -> str: ...
+
+            @overload
+            def good1(x: int) -> int: ...
+            @overload
+            @_dec1
+            def good1(unrelated: int, whatever: str) -> str: ...
+            """,
+            runtime="def good1(x): ...",
+            error=None,
+        )
+        yield Case(
+            stub="""
+            class _dec2:
+                def __init__(self, func: object) -> None: ...
+                def __call__(self, x: str, y: int) -> str: ...
+
+            @overload
+            def good2(x: int) -> str: ...
+            @overload
+            @_dec2
+            def good2(unrelated: int, whatever: str) -> str: ...
+            """,
+            runtime="def good2(x, y=...): ...",
+            error=None,
+        )
+        yield Case(
+            stub="""
+            class _dec3:
+                def __init__(self, func: object) -> None: ...
+                def __call__(self, x: str, y: int) -> str: ...
+
+            @overload
+            def bad(x: int) -> str: ...
+            @overload
+            @_dec3
+            def bad(unrelated: int, whatever: str) -> str: ...
+            """,
+            runtime="def bad(x): ...",
+            error="bad",
+        )
+
     @collect_cases
     def test_property(self) -> Iterator[Case]:
         yield Case(