1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
|
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
|
