sympy/utilities/decorator.py

   1
   2 from sympy.core.add import Add
   3 from sympy.core.sympify import sympify
   4 from sympy.core.relational import Relational
   5
   6 def threaded(**flags):
   7     """Call a function on all elements of composite objects.
   8
   9        This decorator is intended to make it uniformly possible  to apply
  10        functions to all elements of composite or iterable objects or just
  11        to expressions. Currently supported, so called, threadable classes
  12        are Add, Matrix, Relational and all other which implement __iter__
  13        attribute.  If the input expression is not of this type,  then the
  14        function will be applied to this object, as to a singleton.
  15
  16        Functions which will use this decorator must have the following
  17        signature:
  18
  19           @threaded()
  20           def function(expr, *args, **kwargs):
  21
  22        where 'expr' is obligatory threadable object (listed above). Other
  23        arguments, if available, are transferred without any change. As an
  24        example see functions in sympy.simplify module.
  25
  26        By default threading is done on elements of Add instance. To avoid
  27        this behaviour set 'use_add' flag  with False in keyword arguments
  28        (see integrate() for details), e.g:
  29
  30           @threaded(use_add=False)
  31           def function(expr, *args, **kwargs):
  32
  33     """
  34     from sympy.matrices import Matrix
  35
  36     use_add = flags.get('use_add', True)
  37
  38     def threaded_proxy(func):
  39         def threaded_decorator(expr, *args, **kwargs):
  40             if isinstance(expr, Matrix):
  41                 return expr.applyfunc(lambda f: func(f, *args, **kwargs))
  42             elif hasattr(expr, '__iter__'):
  43                 return expr.__class__([ func(f, *args, **kwargs) for f in expr ])
  44             else:
  45                 expr = sympify(expr)
  46
  47                 if isinstance(expr, Relational):
  48                     lhs = func(expr.lhs, *args, **kwargs)
  49                     rhs = func(expr.rhs, *args, **kwargs)
  50
  51                     return Relational(lhs, rhs, expr.rel_op)
  52                 elif expr.is_Add and use_add:
  53                     return Add(*[ func(f, *args, **kwargs) for f in expr.args ])
  54                 else:
  55                     return func(expr, *args, **kwargs)
  56
  57         threaded_decorator.__doc__  = func.__doc__
  58         threaded_decorator.__name__ = func.__name__
  59
  60         return threaded_decorator
  61
  62     return threaded_proxy