Skip to content

Commit

Permalink
Call get_method_hook when methods are used in decorators (#10675)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
pranavrajpal committed Jun 22, 2021
1 parent 08e6ba8 commit e4c2668
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 1 deletion.
9 changes: 8 additions & 1 deletion mypy/checker.py
Expand Up @@ -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
Expand Down
20 changes: 20 additions & 0 deletions test-data/unit/check-custom-plugin.test
Expand Up @@ -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
19 changes: 19 additions & 0 deletions 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

0 comments on commit e4c2668

Please sign in to comment.