symbolic.pyi 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219
  1. from collections.abc import Callable, Mapping
  2. from enum import Enum
  3. from typing import Any, Generic, Literal as L, ParamSpec, Self, TypeAlias, overload
  4. from typing_extensions import TypeVar
  5. __all__ = ["Expr"]
  6. ###
  7. _Tss = ParamSpec("_Tss")
  8. _ExprT = TypeVar("_ExprT", bound=Expr)
  9. _ExprT1 = TypeVar("_ExprT1", bound=Expr)
  10. _ExprT2 = TypeVar("_ExprT2", bound=Expr)
  11. _OpT_co = TypeVar("_OpT_co", bound=Op, default=Op, covariant=True)
  12. _LanguageT_co = TypeVar("_LanguageT_co", bound=Language, default=Language, covariant=True)
  13. _DataT_co = TypeVar("_DataT_co", default=Any, covariant=True)
  14. _LeftT_co = TypeVar("_LeftT_co", default=Any, covariant=True)
  15. _RightT_co = TypeVar("_RightT_co", default=Any, covariant=True)
  16. _RelCOrPy: TypeAlias = L["==", "!=", "<", "<=", ">", ">="]
  17. _RelFortran: TypeAlias = L[".eq.", ".ne.", ".lt.", ".le.", ".gt.", ".ge."]
  18. _ToExpr: TypeAlias = Expr | complex | str
  19. _ToExprN: TypeAlias = _ToExpr | tuple[_ToExprN, ...]
  20. _NestedString: TypeAlias = str | tuple[_NestedString, ...] | list[_NestedString]
  21. ###
  22. class OpError(Exception): ...
  23. class ExprWarning(UserWarning): ...
  24. class Language(Enum):
  25. Python = 0
  26. Fortran = 1
  27. C = 2
  28. class Op(Enum):
  29. INTEGER = 10
  30. REAL = 12
  31. COMPLEX = 15
  32. STRING = 20
  33. ARRAY = 30
  34. SYMBOL = 40
  35. TERNARY = 100
  36. APPLY = 200
  37. INDEXING = 210
  38. CONCAT = 220
  39. RELATIONAL = 300
  40. TERMS = 1_000
  41. FACTORS = 2_000
  42. REF = 3_000
  43. DEREF = 3_001
  44. class RelOp(Enum):
  45. EQ = 1
  46. NE = 2
  47. LT = 3
  48. LE = 4
  49. GT = 5
  50. GE = 6
  51. @overload
  52. @classmethod
  53. def fromstring(cls, s: _RelCOrPy, language: L[Language.C, Language.Python] = ...) -> RelOp: ...
  54. @overload
  55. @classmethod
  56. def fromstring(cls, s: _RelFortran, language: L[Language.Fortran]) -> RelOp: ...
  57. #
  58. @overload
  59. def tostring(self, /, language: L[Language.C, Language.Python] = ...) -> _RelCOrPy: ...
  60. @overload
  61. def tostring(self, /, language: L[Language.Fortran]) -> _RelFortran: ...
  62. class ArithOp(Enum):
  63. POS = 1
  64. NEG = 2
  65. ADD = 3
  66. SUB = 4
  67. MUL = 5
  68. DIV = 6
  69. POW = 7
  70. class Precedence(Enum):
  71. ATOM = 0
  72. POWER = 1
  73. UNARY = 2
  74. PRODUCT = 3
  75. SUM = 4
  76. LT = 6
  77. EQ = 7
  78. LAND = 11
  79. LOR = 12
  80. TERNARY = 13
  81. ASSIGN = 14
  82. TUPLE = 15
  83. NONE = 100
  84. class Expr(Generic[_OpT_co, _DataT_co]):
  85. op: _OpT_co
  86. data: _DataT_co
  87. @staticmethod
  88. def parse(s: str, language: Language = ...) -> Expr: ...
  89. #
  90. def __init__(self, /, op: Op, data: _DataT_co) -> None: ...
  91. #
  92. def __lt__(self, other: Expr, /) -> bool: ...
  93. def __le__(self, other: Expr, /) -> bool: ...
  94. def __gt__(self, other: Expr, /) -> bool: ...
  95. def __ge__(self, other: Expr, /) -> bool: ...
  96. #
  97. def __pos__(self, /) -> Self: ...
  98. def __neg__(self, /) -> Expr: ...
  99. #
  100. def __add__(self, other: Expr, /) -> Expr: ...
  101. def __radd__(self, other: Expr, /) -> Expr: ...
  102. #
  103. def __sub__(self, other: Expr, /) -> Expr: ...
  104. def __rsub__(self, other: Expr, /) -> Expr: ...
  105. #
  106. def __mul__(self, other: Expr, /) -> Expr: ...
  107. def __rmul__(self, other: Expr, /) -> Expr: ...
  108. #
  109. def __pow__(self, other: Expr, /) -> Expr: ...
  110. #
  111. def __truediv__(self, other: Expr, /) -> Expr: ...
  112. def __rtruediv__(self, other: Expr, /) -> Expr: ...
  113. #
  114. def __floordiv__(self, other: Expr, /) -> Expr: ...
  115. def __rfloordiv__(self, other: Expr, /) -> Expr: ...
  116. #
  117. def __call__(
  118. self,
  119. /,
  120. *args: _ToExprN,
  121. **kwargs: _ToExprN,
  122. ) -> Expr[L[Op.APPLY], tuple[Self, tuple[Expr, ...], dict[str, Expr]]]: ...
  123. #
  124. @overload
  125. def __getitem__(self, index: _ExprT | tuple[_ExprT], /) -> Expr[L[Op.INDEXING], tuple[Self, _ExprT]]: ...
  126. @overload
  127. def __getitem__(self, index: _ToExpr | tuple[_ToExpr], /) -> Expr[L[Op.INDEXING], tuple[Self, Expr]]: ...
  128. #
  129. def substitute(self, /, symbols_map: Mapping[Expr, Expr]) -> Expr: ...
  130. #
  131. @overload
  132. def traverse(self, /, visit: Callable[_Tss, None], *args: _Tss.args, **kwargs: _Tss.kwargs) -> Expr: ...
  133. @overload
  134. def traverse(self, /, visit: Callable[_Tss, _ExprT], *args: _Tss.args, **kwargs: _Tss.kwargs) -> _ExprT: ...
  135. #
  136. def contains(self, /, other: Expr) -> bool: ...
  137. #
  138. def symbols(self, /) -> set[Expr]: ...
  139. def polynomial_atoms(self, /) -> set[Expr]: ...
  140. #
  141. def linear_solve(self, /, symbol: Expr) -> tuple[Expr, Expr]: ...
  142. #
  143. def tostring(self, /, parent_precedence: Precedence = ..., language: Language = ...) -> str: ...
  144. class _Pair(Generic[_LeftT_co, _RightT_co]):
  145. left: _LeftT_co
  146. right: _RightT_co
  147. def __init__(self, /, left: _LeftT_co, right: _RightT_co) -> None: ...
  148. #
  149. @overload
  150. def substitute(self: _Pair[_ExprT1, _ExprT2], /, symbols_map: Mapping[Expr, Expr]) -> _Pair[Expr, Expr]: ...
  151. @overload
  152. def substitute(self: _Pair[_ExprT1, object], /, symbols_map: Mapping[Expr, Expr]) -> _Pair[Expr, Any]: ...
  153. @overload
  154. def substitute(self: _Pair[object, _ExprT2], /, symbols_map: Mapping[Expr, Expr]) -> _Pair[Any, Expr]: ...
  155. @overload
  156. def substitute(self, /, symbols_map: Mapping[Expr, Expr]) -> _Pair: ...
  157. class _FromStringWorker(Generic[_LanguageT_co]):
  158. language: _LanguageT_co
  159. original: str | None
  160. quotes_map: dict[str, str]
  161. @overload
  162. def __init__(self: _FromStringWorker[L[Language.C]], /, language: L[Language.C] = ...) -> None: ...
  163. @overload
  164. def __init__(self, /, language: _LanguageT_co) -> None: ...
  165. #
  166. def finalize_string(self, /, s: str) -> str: ...
  167. #
  168. def parse(self, /, inp: str) -> Expr | _Pair: ...
  169. #
  170. @overload
  171. def process(self, /, s: str, context: str = "expr") -> Expr | _Pair: ...
  172. @overload
  173. def process(self, /, s: list[str], context: str = "expr") -> list[Expr | _Pair]: ...
  174. @overload
  175. def process(self, /, s: tuple[str, ...], context: str = "expr") -> tuple[Expr | _Pair, ...]: ...
  176. @overload
  177. def process(self, /, s: _NestedString, context: str = "expr") -> Any: ... # noqa: ANN401