Coverage for /usr/lib/python3/dist-packages/sympy/core/decorators.py: 80%
70 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""
2SymPy core decorators.
4The purpose of this module is to expose decorators without any other
5dependencies, so that they can be easily imported anywhere in sympy/core.
6"""
8from functools import wraps
9from .sympify import SympifyError, sympify
12def _sympifyit(arg, retval=None):
13 """
14 decorator to smartly _sympify function arguments
16 Explanation
17 ===========
19 @_sympifyit('other', NotImplemented)
20 def add(self, other):
21 ...
23 In add, other can be thought of as already being a SymPy object.
25 If it is not, the code is likely to catch an exception, then other will
26 be explicitly _sympified, and the whole code restarted.
28 if _sympify(arg) fails, NotImplemented will be returned
30 See also
31 ========
33 __sympifyit
34 """
35 def deco(func):
36 return __sympifyit(func, arg, retval)
38 return deco
41def __sympifyit(func, arg, retval=None):
42 """Decorator to _sympify `arg` argument for function `func`.
44 Do not use directly -- use _sympifyit instead.
45 """
47 # we support f(a,b) only
48 if not func.__code__.co_argcount:
49 raise LookupError("func not found")
50 # only b is _sympified
51 assert func.__code__.co_varnames[1] == arg
52 if retval is None:
53 @wraps(func)
54 def __sympifyit_wrapper(a, b):
55 return func(a, sympify(b, strict=True))
57 else:
58 @wraps(func)
59 def __sympifyit_wrapper(a, b):
60 try:
61 # If an external class has _op_priority, it knows how to deal
62 # with SymPy objects. Otherwise, it must be converted.
63 if not hasattr(b, '_op_priority'):
64 b = sympify(b, strict=True)
65 return func(a, b)
66 except SympifyError:
67 return retval
69 return __sympifyit_wrapper
72def call_highest_priority(method_name):
73 """A decorator for binary special methods to handle _op_priority.
75 Explanation
76 ===========
78 Binary special methods in Expr and its subclasses use a special attribute
79 '_op_priority' to determine whose special method will be called to
80 handle the operation. In general, the object having the highest value of
81 '_op_priority' will handle the operation. Expr and subclasses that define
82 custom binary special methods (__mul__, etc.) should decorate those
83 methods with this decorator to add the priority logic.
85 The ``method_name`` argument is the name of the method of the other class
86 that will be called. Use this decorator in the following manner::
88 # Call other.__rmul__ if other._op_priority > self._op_priority
89 @call_highest_priority('__rmul__')
90 def __mul__(self, other):
91 ...
93 # Call other.__mul__ if other._op_priority > self._op_priority
94 @call_highest_priority('__mul__')
95 def __rmul__(self, other):
96 ...
97 """
98 def priority_decorator(func):
99 @wraps(func)
100 def binary_op_wrapper(self, other):
101 if hasattr(other, '_op_priority'):
102 if other._op_priority > self._op_priority:
103 f = getattr(other, method_name, None)
104 if f is not None:
105 return f(self)
106 return func(self, other)
107 return binary_op_wrapper
108 return priority_decorator
111def sympify_method_args(cls):
112 '''Decorator for a class with methods that sympify arguments.
114 Explanation
115 ===========
117 The sympify_method_args decorator is to be used with the sympify_return
118 decorator for automatic sympification of method arguments. This is
119 intended for the common idiom of writing a class like :
121 Examples
122 ========
124 >>> from sympy import Basic, SympifyError, S
125 >>> from sympy.core.sympify import _sympify
127 >>> class MyTuple(Basic):
128 ... def __add__(self, other):
129 ... try:
130 ... other = _sympify(other)
131 ... except SympifyError:
132 ... return NotImplemented
133 ... if not isinstance(other, MyTuple):
134 ... return NotImplemented
135 ... return MyTuple(*(self.args + other.args))
137 >>> MyTuple(S(1), S(2)) + MyTuple(S(3), S(4))
138 MyTuple(1, 2, 3, 4)
140 In the above it is important that we return NotImplemented when other is
141 not sympifiable and also when the sympified result is not of the expected
142 type. This allows the MyTuple class to be used cooperatively with other
143 classes that overload __add__ and want to do something else in combination
144 with instance of Tuple.
146 Using this decorator the above can be written as
148 >>> from sympy.core.decorators import sympify_method_args, sympify_return
150 >>> @sympify_method_args
151 ... class MyTuple(Basic):
152 ... @sympify_return([('other', 'MyTuple')], NotImplemented)
153 ... def __add__(self, other):
154 ... return MyTuple(*(self.args + other.args))
156 >>> MyTuple(S(1), S(2)) + MyTuple(S(3), S(4))
157 MyTuple(1, 2, 3, 4)
159 The idea here is that the decorators take care of the boiler-plate code
160 for making this happen in each method that potentially needs to accept
161 unsympified arguments. Then the body of e.g. the __add__ method can be
162 written without needing to worry about calling _sympify or checking the
163 type of the resulting object.
165 The parameters for sympify_return are a list of tuples of the form
166 (parameter_name, expected_type) and the value to return (e.g.
167 NotImplemented). The expected_type parameter can be a type e.g. Tuple or a
168 string 'Tuple'. Using a string is useful for specifying a Type within its
169 class body (as in the above example).
171 Notes: Currently sympify_return only works for methods that take a single
172 argument (not including self). Specifying an expected_type as a string
173 only works for the class in which the method is defined.
174 '''
175 # Extract the wrapped methods from each of the wrapper objects created by
176 # the sympify_return decorator. Doing this here allows us to provide the
177 # cls argument which is used for forward string referencing.
178 for attrname, obj in cls.__dict__.items():
179 if isinstance(obj, _SympifyWrapper):
180 setattr(cls, attrname, obj.make_wrapped(cls))
181 return cls
184def sympify_return(*args):
185 '''Function/method decorator to sympify arguments automatically
187 See the docstring of sympify_method_args for explanation.
188 '''
189 # Store a wrapper object for the decorated method
190 def wrapper(func):
191 return _SympifyWrapper(func, args)
192 return wrapper
195class _SympifyWrapper:
196 '''Internal class used by sympify_return and sympify_method_args'''
198 def __init__(self, func, args):
199 self.func = func
200 self.args = args
202 def make_wrapped(self, cls):
203 func = self.func
204 parameters, retval = self.args
206 # XXX: Handle more than one parameter?
207 [(parameter, expectedcls)] = parameters
209 # Handle forward references to the current class using strings
210 if expectedcls == cls.__name__:
211 expectedcls = cls
213 # Raise RuntimeError since this is a failure at import time and should
214 # not be recoverable.
215 nargs = func.__code__.co_argcount
216 # we support f(a, b) only
217 if nargs != 2:
218 raise RuntimeError('sympify_return can only be used with 2 argument functions')
219 # only b is _sympified
220 if func.__code__.co_varnames[1] != parameter:
221 raise RuntimeError('parameter name mismatch "%s" in %s' %
222 (parameter, func.__name__))
224 @wraps(func)
225 def _func(self, other):
226 # XXX: The check for _op_priority here should be removed. It is
227 # needed to stop mutable matrices from being sympified to
228 # immutable matrices which breaks things in quantum...
229 if not hasattr(other, '_op_priority'):
230 try:
231 other = sympify(other, strict=True)
232 except SympifyError:
233 return retval
234 if not isinstance(other, expectedcls):
235 return retval
236 return func(self, other)
238 return _func