support Callable / callable Protocols in suggest decorator unwarpping

Author: asottile

mypy/suggestions.py

@@ -230,6 +230,18 @@ def is_implicit_any(typ: Type) -> bool:
     return isinstance(typ, AnyType) and not is_explicit_any(typ)
 
 
+def _arg_accepts_function(typ: ProperType) -> bool:
+    return (
+        # TypeVar / Callable
+        isinstance(typ, (TypeVarType, CallableType))
+        or
+        # Protocol with __call__
+        isinstance(typ, Instance)
+        and typ.type.is_protocol
+        and isinstance(typ.type.get_method("__call__"), FuncDef)
+    )
+
+
 class SuggestionEngine:
     """Engine for finding call sites and suggesting signatures."""
 
@@ -659,7 +671,7 @@ def extract_from_decorator(self, node: Decorator) -> FuncDef | None:
             for ct in typ.items:
                 if not (
                     len(ct.arg_types) == 1
-                    and isinstance(ct.arg_types[0], TypeVarType)
+                    and _arg_accepts_function(get_proper_type(ct.arg_types[0]))
                     and ct.arg_types[0] == ct.ret_type
                 ):
                     return None

test-data/unit/fine-grained-suggest.test

@@ -651,6 +651,38 @@ foo3('hello hello')
 (str) -> str
 ==
 
+[case testSuggestInferFuncDecorator6]
+# suggest: foo.f
+[file foo.py]
+from __future__ import annotations
+
+from typing import Callable, Protocol, TypeVar
+from typing_extensions import ParamSpec
+
+P = ParamSpec('P')
+R = TypeVar('R')
+R_co = TypeVar('R_co', covariant=True)
+
+class Proto(Protocol[P, R_co]):
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R_co: ...
+
+def dec1(f: Callable[P, R]) -> Callable[P, R]: ...
+def dec2(f: Callable[..., R]) -> Callable[..., R]: ...
+def dec3(f: Proto[P, R_co]) -> Proto[P, R_co]: ...
+
+@dec1
+@dec2
+@dec3
+def f(x):
+    return x
+
+f('hi')
+
+[builtins fixtures/isinstancelist.pyi]
+[out]
+(str) -> str
+==
+
 [case testSuggestFlexAny1]
 # suggest: --flex-any=0.4 m.foo
 # suggest: --flex-any=0.7 m.foo