symbolic.pyi 5.9 KB

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