setexpr.py 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. from sympy.core import Expr
  2. from sympy.core.decorators import call_highest_priority, _sympifyit
  3. from .fancysets import ImageSet
  4. from .sets import set_add, set_sub, set_mul, set_div, set_pow, set_function
  5. class SetExpr(Expr):
  6. """An expression that can take on values of a set.
  7. Examples
  8. ========
  9. >>> from sympy import Interval, FiniteSet
  10. >>> from sympy.sets.setexpr import SetExpr
  11. >>> a = SetExpr(Interval(0, 5))
  12. >>> b = SetExpr(FiniteSet(1, 10))
  13. >>> (a + b).set
  14. Union(Interval(1, 6), Interval(10, 15))
  15. >>> (2*a + b).set
  16. Interval(1, 20)
  17. """
  18. _op_priority = 11.0
  19. def __new__(cls, setarg):
  20. return Expr.__new__(cls, setarg)
  21. set = property(lambda self: self.args[0])
  22. def _latex(self, printer):
  23. return r"SetExpr\left({}\right)".format(printer._print(self.set))
  24. @_sympifyit('other', NotImplemented)
  25. @call_highest_priority('__radd__')
  26. def __add__(self, other):
  27. return _setexpr_apply_operation(set_add, self, other)
  28. @_sympifyit('other', NotImplemented)
  29. @call_highest_priority('__add__')
  30. def __radd__(self, other):
  31. return _setexpr_apply_operation(set_add, other, self)
  32. @_sympifyit('other', NotImplemented)
  33. @call_highest_priority('__rmul__')
  34. def __mul__(self, other):
  35. return _setexpr_apply_operation(set_mul, self, other)
  36. @_sympifyit('other', NotImplemented)
  37. @call_highest_priority('__mul__')
  38. def __rmul__(self, other):
  39. return _setexpr_apply_operation(set_mul, other, self)
  40. @_sympifyit('other', NotImplemented)
  41. @call_highest_priority('__rsub__')
  42. def __sub__(self, other):
  43. return _setexpr_apply_operation(set_sub, self, other)
  44. @_sympifyit('other', NotImplemented)
  45. @call_highest_priority('__sub__')
  46. def __rsub__(self, other):
  47. return _setexpr_apply_operation(set_sub, other, self)
  48. @_sympifyit('other', NotImplemented)
  49. @call_highest_priority('__rpow__')
  50. def __pow__(self, other):
  51. return _setexpr_apply_operation(set_pow, self, other)
  52. @_sympifyit('other', NotImplemented)
  53. @call_highest_priority('__pow__')
  54. def __rpow__(self, other):
  55. return _setexpr_apply_operation(set_pow, other, self)
  56. @_sympifyit('other', NotImplemented)
  57. @call_highest_priority('__rtruediv__')
  58. def __truediv__(self, other):
  59. return _setexpr_apply_operation(set_div, self, other)
  60. @_sympifyit('other', NotImplemented)
  61. @call_highest_priority('__truediv__')
  62. def __rtruediv__(self, other):
  63. return _setexpr_apply_operation(set_div, other, self)
  64. def _eval_func(self, func):
  65. # TODO: this could be implemented straight into `imageset`:
  66. res = set_function(func, self.set)
  67. if res is None:
  68. return SetExpr(ImageSet(func, self.set))
  69. return SetExpr(res)
  70. def _setexpr_apply_operation(op, x, y):
  71. if isinstance(x, SetExpr):
  72. x = x.set
  73. if isinstance(y, SetExpr):
  74. y = y.set
  75. out = op(x, y)
  76. return SetExpr(out)