Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Call get_method_hook when methods are used in decorators #10675

Merged
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion mypy/checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -3592,10 +3592,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
41 changes: 41 additions & 0 deletions test-data/unit/check-custom-plugin.test
Original file line number Diff line number Diff line change
Expand Up @@ -783,3 +783,44 @@ 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

[case testPluginMethodHookMethodAssignedToVariableInDecorator-xfail]
pranavrajpal marked this conversation as resolved.
Show resolved Hide resolved
# 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()

meth = b.a
@meth
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
Original file line number Diff line number Diff line change
@@ -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