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
56
57
58
59
60
61
62
63
64
65
66
67
|
from mypy.plugin import Plugin
from mypy.plugins.common import add_attribute_to_class
from mypy.types import NoneType, UnionType
class SympyPlugin(Plugin):
def get_base_class_hook(self, fullname: str):
# TODO: This apparently never worked
if fullname == "sympy.core.basic.Basic":
return add_assumptions
return None
def get_attribute_hook(self, fullname: str):
if fullname == "sympy.core.basic.Basic.free_symbols":
return lambda ctx: ctx.api.named_generic_type(
"builtins.set", [ctx.api.named_type("sympy.Symbol")]
)
return None
def add_assumptions(ctx) -> None:
# Generated by list(sys.modules['sympy.core.assumptions']._assume_defined)
# (do not import sympy to speedup mypy plugin load time)
assumptions = [
"hermitian",
"prime",
"noninteger",
"negative",
"antihermitian",
"infinite",
"finite",
"irrational",
"extended_positive",
"nonpositive",
"odd",
"algebraic",
"integer",
"rational",
"extended_real",
"nonnegative",
"transcendental",
"extended_nonzero",
"extended_negative",
"composite",
"complex",
"imaginary",
"nonzero",
"zero",
"even",
"positive",
"polar",
"extended_nonpositive",
"extended_nonnegative",
"real",
"commutative",
]
for a in assumptions:
add_attribute_to_class(
ctx.api,
ctx.cls,
f"is_{a}",
UnionType([ctx.api.named_type("builtins.bool"), NoneType()]),
)
def plugin(version: str):
return SympyPlugin
|
