containers.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415
  1. """Module for SymPy containers
  2. (SymPy objects that store other SymPy objects)
  3. The containers implemented in this module are subclassed to Basic.
  4. They are supposed to work seamlessly within the SymPy framework.
  5. """
  6. from __future__ import annotations
  7. from collections import OrderedDict
  8. from collections.abc import MutableSet
  9. from typing import Any, Callable
  10. from .basic import Basic
  11. from .sorting import default_sort_key, ordered
  12. from .sympify import _sympify, sympify, _sympy_converter, SympifyError
  13. from sympy.core.kind import Kind
  14. from sympy.utilities.iterables import iterable
  15. from sympy.utilities.misc import as_int
  16. class Tuple(Basic):
  17. """
  18. Wrapper around the builtin tuple object.
  19. Explanation
  20. ===========
  21. The Tuple is a subclass of Basic, so that it works well in the
  22. SymPy framework. The wrapped tuple is available as self.args, but
  23. you can also access elements or slices with [:] syntax.
  24. Parameters
  25. ==========
  26. sympify : bool
  27. If ``False``, ``sympify`` is not called on ``args``. This
  28. can be used for speedups for very large tuples where the
  29. elements are known to already be SymPy objects.
  30. Examples
  31. ========
  32. >>> from sympy import Tuple, symbols
  33. >>> a, b, c, d = symbols('a b c d')
  34. >>> Tuple(a, b, c)[1:]
  35. (b, c)
  36. >>> Tuple(a, b, c).subs(a, d)
  37. (d, b, c)
  38. """
  39. def __new__(cls, *args, **kwargs):
  40. if kwargs.get('sympify', True):
  41. args = (sympify(arg) for arg in args)
  42. obj = Basic.__new__(cls, *args)
  43. return obj
  44. def __getitem__(self, i):
  45. if isinstance(i, slice):
  46. indices = i.indices(len(self))
  47. return Tuple(*(self.args[j] for j in range(*indices)))
  48. return self.args[i]
  49. def __len__(self):
  50. return len(self.args)
  51. def __contains__(self, item):
  52. return item in self.args
  53. def __iter__(self):
  54. return iter(self.args)
  55. def __add__(self, other):
  56. if isinstance(other, Tuple):
  57. return Tuple(*(self.args + other.args))
  58. elif isinstance(other, tuple):
  59. return Tuple(*(self.args + other))
  60. else:
  61. return NotImplemented
  62. def __radd__(self, other):
  63. if isinstance(other, Tuple):
  64. return Tuple(*(other.args + self.args))
  65. elif isinstance(other, tuple):
  66. return Tuple(*(other + self.args))
  67. else:
  68. return NotImplemented
  69. def __mul__(self, other):
  70. try:
  71. n = as_int(other)
  72. except ValueError:
  73. raise TypeError("Can't multiply sequence by non-integer of type '%s'" % type(other))
  74. return self.func(*(self.args*n))
  75. __rmul__ = __mul__
  76. def __eq__(self, other):
  77. if isinstance(other, Basic):
  78. return super().__eq__(other)
  79. return self.args == other
  80. def __ne__(self, other):
  81. if isinstance(other, Basic):
  82. return super().__ne__(other)
  83. return self.args != other
  84. def __hash__(self):
  85. return hash(self.args)
  86. def _to_mpmath(self, prec):
  87. return tuple(a._to_mpmath(prec) for a in self.args)
  88. def __lt__(self, other):
  89. return _sympify(self.args < other.args)
  90. def __le__(self, other):
  91. return _sympify(self.args <= other.args)
  92. # XXX: Basic defines count() as something different, so we can't
  93. # redefine it here. Originally this lead to cse() test failure.
  94. def tuple_count(self, value) -> int:
  95. """Return number of occurrences of value."""
  96. return self.args.count(value)
  97. def index(self, value, start=None, stop=None):
  98. """Searches and returns the first index of the value."""
  99. # XXX: One would expect:
  100. #
  101. # return self.args.index(value, start, stop)
  102. #
  103. # here. Any trouble with that? Yes:
  104. #
  105. # >>> (1,).index(1, None, None)
  106. # Traceback (most recent call last):
  107. # File "<stdin>", line 1, in <module>
  108. # TypeError: slice indices must be integers or None or have an __index__ method
  109. #
  110. # See: http://bugs.python.org/issue13340
  111. if start is None and stop is None:
  112. return self.args.index(value)
  113. elif stop is None:
  114. return self.args.index(value, start)
  115. else:
  116. return self.args.index(value, start, stop)
  117. @property
  118. def kind(self):
  119. """
  120. The kind of a Tuple instance.
  121. The kind of a Tuple is always of :class:`TupleKind` but
  122. parametrised by the number of elements and the kind of each element.
  123. Examples
  124. ========
  125. >>> from sympy import Tuple, Matrix
  126. >>> Tuple(1, 2).kind
  127. TupleKind(NumberKind, NumberKind)
  128. >>> Tuple(Matrix([1, 2]), 1).kind
  129. TupleKind(MatrixKind(NumberKind), NumberKind)
  130. >>> Tuple(1, 2).kind.element_kind
  131. (NumberKind, NumberKind)
  132. See Also
  133. ========
  134. sympy.matrices.kind.MatrixKind
  135. sympy.core.kind.NumberKind
  136. """
  137. return TupleKind(*(i.kind for i in self.args))
  138. _sympy_converter[tuple] = lambda tup: Tuple(*tup)
  139. def tuple_wrapper(method):
  140. """
  141. Decorator that converts any tuple in the function arguments into a Tuple.
  142. Explanation
  143. ===========
  144. The motivation for this is to provide simple user interfaces. The user can
  145. call a function with regular tuples in the argument, and the wrapper will
  146. convert them to Tuples before handing them to the function.
  147. Explanation
  148. ===========
  149. >>> from sympy.core.containers import tuple_wrapper
  150. >>> def f(*args):
  151. ... return args
  152. >>> g = tuple_wrapper(f)
  153. The decorated function g sees only the Tuple argument:
  154. >>> g(0, (1, 2), 3)
  155. (0, (1, 2), 3)
  156. """
  157. def wrap_tuples(*args, **kw_args):
  158. newargs = []
  159. for arg in args:
  160. if isinstance(arg, tuple):
  161. newargs.append(Tuple(*arg))
  162. else:
  163. newargs.append(arg)
  164. return method(*newargs, **kw_args)
  165. return wrap_tuples
  166. class Dict(Basic):
  167. """
  168. Wrapper around the builtin dict object.
  169. Explanation
  170. ===========
  171. The Dict is a subclass of Basic, so that it works well in the
  172. SymPy framework. Because it is immutable, it may be included
  173. in sets, but its values must all be given at instantiation and
  174. cannot be changed afterwards. Otherwise it behaves identically
  175. to the Python dict.
  176. Examples
  177. ========
  178. >>> from sympy import Dict, Symbol
  179. >>> D = Dict({1: 'one', 2: 'two'})
  180. >>> for key in D:
  181. ... if key == 1:
  182. ... print('%s %s' % (key, D[key]))
  183. 1 one
  184. The args are sympified so the 1 and 2 are Integers and the values
  185. are Symbols. Queries automatically sympify args so the following work:
  186. >>> 1 in D
  187. True
  188. >>> D.has(Symbol('one')) # searches keys and values
  189. True
  190. >>> 'one' in D # not in the keys
  191. False
  192. >>> D[1]
  193. one
  194. """
  195. elements: frozenset[Tuple]
  196. _dict: dict[Basic, Basic]
  197. def __new__(cls, *args):
  198. if len(args) == 1 and isinstance(args[0], (dict, Dict)):
  199. items = [Tuple(k, v) for k, v in args[0].items()]
  200. elif iterable(args) and all(len(arg) == 2 for arg in args):
  201. items = [Tuple(k, v) for k, v in args]
  202. else:
  203. raise TypeError('Pass Dict args as Dict((k1, v1), ...) or Dict({k1: v1, ...})')
  204. elements = frozenset(items)
  205. obj = Basic.__new__(cls, *ordered(items))
  206. obj.elements = elements
  207. obj._dict = dict(items) # In case Tuple decides it wants to sympify
  208. return obj
  209. def __getitem__(self, key):
  210. """x.__getitem__(y) <==> x[y]"""
  211. try:
  212. key = _sympify(key)
  213. except SympifyError:
  214. raise KeyError(key)
  215. return self._dict[key]
  216. def __setitem__(self, key, value):
  217. raise NotImplementedError("SymPy Dicts are Immutable")
  218. def items(self):
  219. '''Returns a set-like object providing a view on dict's items.
  220. '''
  221. return self._dict.items()
  222. def keys(self):
  223. '''Returns the list of the dict's keys.'''
  224. return self._dict.keys()
  225. def values(self):
  226. '''Returns the list of the dict's values.'''
  227. return self._dict.values()
  228. def __iter__(self):
  229. '''x.__iter__() <==> iter(x)'''
  230. return iter(self._dict)
  231. def __len__(self):
  232. '''x.__len__() <==> len(x)'''
  233. return self._dict.__len__()
  234. def get(self, key, default=None):
  235. '''Returns the value for key if the key is in the dictionary.'''
  236. try:
  237. key = _sympify(key)
  238. except SympifyError:
  239. return default
  240. return self._dict.get(key, default)
  241. def __contains__(self, key):
  242. '''D.__contains__(k) -> True if D has a key k, else False'''
  243. try:
  244. key = _sympify(key)
  245. except SympifyError:
  246. return False
  247. return key in self._dict
  248. def __lt__(self, other):
  249. return _sympify(self.args < other.args)
  250. @property
  251. def _sorted_args(self):
  252. return tuple(sorted(self.args, key=default_sort_key))
  253. def __eq__(self, other):
  254. if isinstance(other, dict):
  255. return self == Dict(other)
  256. return super().__eq__(other)
  257. __hash__ : Callable[[Basic], Any] = Basic.__hash__
  258. # this handles dict, defaultdict, OrderedDict
  259. _sympy_converter[dict] = lambda d: Dict(*d.items())
  260. class OrderedSet(MutableSet):
  261. def __init__(self, iterable=None):
  262. if iterable:
  263. self.map = OrderedDict((item, None) for item in iterable)
  264. else:
  265. self.map = OrderedDict()
  266. def __len__(self):
  267. return len(self.map)
  268. def __contains__(self, key):
  269. return key in self.map
  270. def add(self, key):
  271. self.map[key] = None
  272. def discard(self, key):
  273. self.map.pop(key)
  274. def pop(self, last=True):
  275. return self.map.popitem(last=last)[0]
  276. def __iter__(self):
  277. yield from self.map.keys()
  278. def __repr__(self):
  279. if not self.map:
  280. return '%s()' % (self.__class__.__name__,)
  281. return '%s(%r)' % (self.__class__.__name__, list(self.map.keys()))
  282. def intersection(self, other):
  283. return self.__class__([val for val in self if val in other])
  284. def difference(self, other):
  285. return self.__class__([val for val in self if val not in other])
  286. def update(self, iterable):
  287. for val in iterable:
  288. self.add(val)
  289. class TupleKind(Kind):
  290. """
  291. TupleKind is a subclass of Kind, which is used to define Kind of ``Tuple``.
  292. Parameters of TupleKind will be kinds of all the arguments in Tuples, for
  293. example
  294. Parameters
  295. ==========
  296. args : tuple(element_kind)
  297. element_kind is kind of element.
  298. args is tuple of kinds of element
  299. Examples
  300. ========
  301. >>> from sympy import Tuple
  302. >>> Tuple(1, 2).kind
  303. TupleKind(NumberKind, NumberKind)
  304. >>> Tuple(1, 2).kind.element_kind
  305. (NumberKind, NumberKind)
  306. See Also
  307. ========
  308. sympy.core.kind.NumberKind
  309. MatrixKind
  310. sympy.sets.sets.SetKind
  311. """
  312. def __new__(cls, *args):
  313. obj = super().__new__(cls, *args)
  314. obj.element_kind = args
  315. return obj
  316. def __repr__(self):
  317. return "TupleKind{}".format(self.element_kind)