Call get_method_hook when methods are used in decorators (#10675)

pranavrajpal · web-flow · commit e4c2668e29bf · 2021-06-21T17:08:02.000-07:00
When a method is used in a decorator, pass the full name to plugins through the get_method_hook function, just like we would normally do for other places that methods got called. Previously, in these situations, we would call get_function_hook instead with '{method_name} of {class_name}', which loses information about which module the class is located in, and gives information about a method to the hook meant for functions. Currently this is limited to situations where the method is accessed in the decorator itself, so @obj.meth def f() -> None: ... would work but method = obj.meth @method def f() -> None: ... probably wouldn't work.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -3596,10 +3596,17 @@ def visit_decorator(self, e: Decorator) -> None:
             fullname = None
             if isinstance(d, RefExpr):
                 fullname = d.fullname
+            # if this is a expression like @b.a where b is an object, get the type of b
+            # so we can pass it the method hook in the plugins
+            object_type = None  # type: Optional[Type]
+            if fullname is None and isinstance(d, MemberExpr) and d.expr in self.type_map:
+                object_type = self.type_map[d.expr]
+                fullname = self.expr_checker.method_fullname(object_type, d.name)
             self.check_for_untyped_decorator(e.func, dec, d)
             sig, t2 = self.expr_checker.check_call(dec, [temp],
                                                    [nodes.ARG_POS], e,
-                                                   callable_name=fullname)
+                                                   callable_name=fullname,
+                                                   object_type=object_type)
         self.check_untyped_after_decorator(sig, e.func)
         sig = set_callable_name(sig, e.func)
         e.var.type = sig
diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
@@ -783,3 +783,23 @@ reveal_type(dynamic_signature(1))  # N: Revealed type is "builtins.int"
 [file mypy.ini]
 \[mypy]
 plugins=<ROOT>/test-data/unit/plugins/function_sig_hook.py
+
+[case testPluginCalledCorrectlyWhenMethodInDecorator]
+# flags: --config-file tmp/mypy.ini
+from typing import TypeVar, Callable
+
+T = TypeVar('T')
+class Foo:
+    def a(self, x: Callable[[], T]) -> Callable[[], T]: ...
+
+b = Foo()
+
+@b.a
+def f() -> None:
+    pass
+
+reveal_type(f())  # N: Revealed type is "builtins.str"
+
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/method_in_decorator.py
diff --git a/test-data/unit/plugins/method_in_decorator.py b/test-data/unit/plugins/method_in_decorator.py
@@ -0,0 +1,19 @@
+from mypy.types import CallableType, Type
+from typing import Callable, Optional
+from mypy.plugin import MethodContext, Plugin
+
+
+class MethodDecoratorPlugin(Plugin):
+    def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], Type]]:
+        if 'Foo.a' in fullname:
+            return method_decorator_callback
+        return None
+
+def method_decorator_callback(ctx: MethodContext) -> Type:
+    if isinstance(ctx.default_return_type, CallableType):
+        str_type = ctx.api.named_generic_type('builtins.str', [])
+        return ctx.default_return_type.copy_modified(ret_type=str_type)
+    return ctx.default_return_type
+
+def plugin(version):
+    return MethodDecoratorPlugin
