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