dmypy suggest can now suggest through contextmanager-based decorators

Author: asottile

Diff for: mypy/suggestions.py

@@ -52,14 +52,15 @@
     SymbolNode,
     SymbolTable,
     TypeInfo,
+    Var,
     reverse_builtin_aliases,
 )
 from mypy.options import Options
 from mypy.plugin import FunctionContext, MethodContext, Plugin
 from mypy.server.update import FineGrainedBuildManager
 from mypy.state import state
 from mypy.traverser import TraverserVisitor
-from mypy.typeops import make_simplified_union
+from mypy.typeops import bind_self, make_simplified_union
 from mypy.types import (
     AnyType,
     CallableType,
@@ -638,15 +639,20 @@ def find_node_by_file_and_line(self, file: str, line: int) -> tuple[str, SymbolN
     def extract_from_decorator(self, node: Decorator) -> FuncDef | None:
         for dec in node.decorators:
             typ = None
-            if isinstance(dec, RefExpr) and isinstance(dec.node, FuncDef):
-                typ = dec.node.type
+            if isinstance(dec, RefExpr) and isinstance(dec.node, (Var, FuncDef)):
+                typ = get_proper_type(dec.node.type)
             elif (
                 isinstance(dec, CallExpr)
                 and isinstance(dec.callee, RefExpr)
-                and isinstance(dec.callee.node, FuncDef)
-                and isinstance(dec.callee.node.type, CallableType)
+                and isinstance(dec.callee.node, (Decorator, FuncDef, Var))
+                and isinstance((call_tp := get_proper_type(dec.callee.node.type)), CallableType)
             ):
-                typ = get_proper_type(dec.callee.node.type.ret_type)
+                typ = get_proper_type(call_tp.ret_type)
+
+            if isinstance(typ, Instance):
+                call_method = typ.type.get_method("__call__")
+                if isinstance(call_method, FuncDef) and isinstance(call_method.type, FunctionLike):
+                    typ = bind_self(call_method.type, None)
 
             if not isinstance(typ, FunctionLike):
                 return None

Diff for: test-data/unit/fine-grained-suggest.test

@@ -602,6 +602,55 @@ def bar() -> None:
 (str) -> str
 ==
 
+[case testSuggestInferFuncDecorator5]
+# suggest: foo.foo1
+# suggest: foo.foo2
+# suggest: foo.foo3
+[file foo.py]
+from __future__ import annotations
+
+from typing import TypeVar, Generator, Callable
+
+F = TypeVar('F')
+
+# simplified `@contextmanager
+class _impl:
+    def __call__(self, f: F) -> F: return f
+def contextmanager(gen: Callable[[], Generator[None, None, None]]) -> Callable[[], _impl]: return _impl
+
+@contextmanager
+def gen() -> Generator[None, None, None]:
+    yield
+
+@gen()
+def foo1(x):
+    return x
+
+foo1('hi')
+
+inst = gen()
+
+@inst
+def foo2(x):
+    return x
+
+foo2('hello')
+
+ref = gen
+
+@ref()
+def foo3(x):
+    return x
+
+foo3('hello hello')
+
+[builtins fixtures/isinstancelist.pyi]
+[out]
+(str) -> str
+(str) -> str
+(str) -> str
+==
+
 [case testSuggestFlexAny1]
 # suggest: --flex-any=0.4 m.foo
 # suggest: --flex-any=0.7 m.foo