| from mypy.plugin import MethodContext, Plugin | |
| from mypy.types import Instance, Type | |
| class CallableInstancePlugin(Plugin): | |
| def get_function_hook(self, fullname): | |
| assert not fullname.endswith(' of Foo') | |
| def get_method_hook(self, fullname): | |
| # Ensure that all names are fully qualified | |
| assert not fullname.endswith(' of Foo') | |
| if fullname == '__main__.Class.__call__': | |
| return my_hook | |
| return None | |
| def my_hook(ctx: MethodContext) -> Type: | |
| if isinstance(ctx.type, Instance) and len(ctx.type.args) == 1: | |
| return ctx.type.args[0] | |
| return ctx.default_return_type | |
| def plugin(version): | |
| return CallableInstancePlugin |