containers.py 11 KB

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