from __future__ import annotations

from typing import Callable

from mypy.plugin import MethodContext, Plugin

from mypy.types import Instance, Type

class CallableInstancePlugin(Plugin):

def get_function_hook(self, fullname: str) -> None:

assert not fullname.endswith(" of Foo")

def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | None:

# 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: str) -> type[CallableInstancePlugin]:

return CallableInstancePlugin