pytorch

1
from mypy.plugin import Plugin
2
from mypy.plugins.common import add_attribute_to_class
3
from mypy.types import NoneType, UnionType
4

5

6
class SympyPlugin(Plugin):
7
    def get_base_class_hook(self, fullname: str):
8
        if fullname == "sympy.core.basic.Basic":
9
            return add_assumptions
10
        return None
11

12

13
def add_assumptions(ctx) -> None:
14
    # Generated by list(sys.modules['sympy.core.assumptions']._assume_defined)
15
    # (do not import sympy to speedup mypy plugin load time)
16
    assumptions = [
17
        "hermitian",
18
        "prime",
19
        "noninteger",
20
        "negative",
21
        "antihermitian",
22
        "infinite",
23
        "finite",
24
        "irrational",
25
        "extended_positive",
26
        "nonpositive",
27
        "odd",
28
        "algebraic",
29
        "integer",
30
        "rational",
31
        "extended_real",
32
        "nonnegative",
33
        "transcendental",
34
        "extended_nonzero",
35
        "extended_negative",
36
        "composite",
37
        "complex",
38
        "imaginary",
39
        "nonzero",
40
        "zero",
41
        "even",
42
        "positive",
43
        "polar",
44
        "extended_nonpositive",
45
        "extended_nonnegative",
46
        "real",
47
        "commutative",
48
    ]
49
    for a in assumptions:
50
        add_attribute_to_class(
51
            ctx.api,
52
            ctx.cls,
53
            f"is_{a}",
54
            UnionType([ctx.api.named_type("builtins.bool"), NoneType()]),
55
        )
56

57

58
def plugin(version: str):
59
    return SympyPlugin
60