from __future__ import annotations

from typing import Callable

from mypy.plugin import CheckerPluginInterface, MethodContext, MethodSigContext, Plugin

from mypy.types import CallableType, Instance, Type, get_proper_type

class MethodPlugin(Plugin):

def get_method_signature_hook(

self, fullname: str

) -> Callable[[MethodSigContext], CallableType] | None:

if fullname.startswith("__main__.Foo."):

return my_meth_sig_hook

return None

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

if fullname.startswith("__main__.Bar."):

return my_meth_hook

return None

def _str_to_int(api: CheckerPluginInterface, typ: Type) -> Type:

typ = get_proper_type(typ)

if isinstance(typ, Instance):

if typ.type.fullname == "builtins.str":

return api.named_generic_type("builtins.int", [])

elif typ.args:

return typ.copy_modified(args=[_str_to_int(api, t) for t in typ.args])

return typ

def _float_to_int(api: CheckerPluginInterface, typ: Type) -> Type:

typ = get_proper_type(typ)

if isinstance(typ, Instance):

if typ.type.fullname == "builtins.float":

return api.named_generic_type("builtins.int", [])

elif typ.args:

return typ.copy_modified(args=[_float_to_int(api, t) for t in typ.args])

return typ

def my_meth_sig_hook(ctx: MethodSigContext) -> CallableType:

return ctx.default_signature.copy_modified(

arg_types=[_str_to_int(ctx.api, t) for t in ctx.default_signature.arg_types],

ret_type=_str_to_int(ctx.api, ctx.default_signature.ret_type),

)

def my_meth_hook(ctx: MethodContext) -> Type:

return _float_to_int(ctx.api, ctx.default_return_type)

def plugin(version: str) -> type[MethodPlugin]:

return MethodPlugin