| 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 |
