| 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): |
| if fullname == "sympy.core.basic.Basic": |
| return add_assumptions |
| 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 |
