File: method_sig_hook.py

package info (click to toggle)
mypy 1.19.1-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 22,412 kB
  • sloc: python: 114,754; ansic: 13,343; cpp: 11,380; makefile: 257; sh: 28
file content (41 lines) | stat: -rw-r--r-- 1,270 bytes parent folder | download | duplicates (2)
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
from __future__ import annotations

from typing import Callable

from mypy.plugin import CheckerPluginInterface, MethodSigContext, Plugin
from mypy.types import CallableType, Instance, Type, get_proper_type


class MethodSigPlugin(Plugin):
    def get_method_signature_hook(
        self, fullname: str
    ) -> Callable[[MethodSigContext], CallableType] | None:
        # Ensure that all names are fully qualified
        assert not fullname.endswith(" of Foo")

        if fullname.startswith("__main__.Foo."):
            return my_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 my_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 plugin(version: str) -> type[MethodSigPlugin]:
    return MethodSigPlugin