diff --git a/mypy/suggestions.py b/mypy/suggestions.py index 16e630bf8c6e..f27ad7cdb637 100644 --- a/mypy/suggestions.py +++ b/mypy/suggestions.py @@ -52,6 +52,7 @@ SymbolNode, SymbolTable, TypeInfo, + Var, reverse_builtin_aliases, ) from mypy.options import Options @@ -59,7 +60,7 @@ 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 --git a/test-data/unit/fine-grained-suggest.test b/test-data/unit/fine-grained-suggest.test index 0ed3be4055ea..2539886229cf 100644 --- a/test-data/unit/fine-grained-suggest.test +++ b/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