Coverage for chebpy/core/decorators.py: 99%
73 statements
« prev ^ index » next coverage.py v7.10.2, created at 2025-08-07 10:30 +0000
1"""Decorator functions for the ChebPy package.
3This module provides various decorators used throughout the ChebPy package to
4implement common functionality such as caching, handling empty objects,
5pre- and post-processing of function inputs/outputs, and type conversion.
6These decorators help reduce code duplication and ensure consistent behavior
7across the package.
8"""
10from functools import wraps
12import numpy as np
15def cache(f: callable) -> callable:
16 """Object method output caching mechanism.
18 This decorator caches the output of zero-argument methods to speed up repeated
19 execution of relatively expensive operations such as .roots(). Cached computations
20 are stored in a dictionary called _cache which is bound to self using keys
21 corresponding to the method name.
23 Args:
24 f (callable): The method to be cached. Must be a zero-argument method.
26 Returns:
27 callable: A wrapped version of the method that implements caching.
29 Note:
30 Can be used in principle on arbitrary objects.
31 """
33 # TODO: look into replacing this with one of the functools cache decorators
34 @wraps(f)
35 def wrapper(self):
36 try:
37 # f has been executed previously
38 out = self._cache[f.__name__]
39 except AttributeError:
40 # f has not been executed previously and self._cache does not exist
41 self._cache = {}
42 out = self._cache[f.__name__] = f(self)
43 except KeyError: # pragma: no cover
44 # f has not been executed previously, but self._cache exists
45 out = self._cache[f.__name__] = f(self)
46 return out
48 return wrapper
51def self_empty(resultif=None) -> callable:
52 """Factory method to produce a decorator for handling empty objects.
54 This factory creates a decorator that checks whether the object whose method
55 is being wrapped is empty. If the object is empty, it returns either the supplied
56 resultif value or a copy of the object. Otherwise, it executes the wrapped method.
58 Args:
59 resultif: Value to return when the object is empty. If None, returns a copy
60 of the object instead.
62 Returns:
63 callable: A decorator function that implements the empty-checking logic.
65 Note:
66 This decorator is primarily used in chebtech.py.
67 """
69 # TODO: add unit test for this
70 def decorator(f):
71 @wraps(f)
72 def wrapper(self, *args, **kwargs):
73 if self.isempty:
74 if resultif is not None:
75 return resultif
76 else:
77 return self.copy()
78 else:
79 return f(self, *args, **kwargs)
81 return wrapper
83 return decorator
86def preandpostprocess(f: callable) -> callable:
87 """Decorator for pre- and post-processing tasks common to bary and clenshaw.
89 This decorator handles several edge cases for functions like bary and clenshaw:
90 - Empty arrays in input arguments
91 - Constant functions
92 - NaN values in coefficients
93 - Scalar vs. array inputs
95 Args:
96 f (callable): The function to be wrapped.
98 Returns:
99 callable: A wrapped version of the function with pre- and post-processing.
100 """
102 @wraps(f)
103 def thewrapper(*args, **kwargs):
104 xx, akfk = args[:2]
105 # are any of the first two arguments empty arrays?
106 if (np.asarray(xx).size == 0) | (np.asarray(akfk).size == 0):
107 return np.array([])
108 # is the function constant?
109 elif akfk.size == 1:
110 if np.isscalar(xx):
111 return akfk[0]
112 else:
113 return akfk * np.ones(xx.size)
114 # are there any NaNs in the second argument?
115 elif np.any(np.isnan(akfk)):
116 return np.nan * np.ones(xx.size)
117 # convert first argument to an array if it is a scalar and then
118 # return the first (and only) element of the result if so
119 else:
120 args = list(args)
121 args[0] = np.array([xx]) if np.isscalar(xx) else args[0]
122 out = f(*args, **kwargs)
123 return out[0] if np.isscalar(xx) else out
125 return thewrapper
128def float_argument(f: callable) -> callable:
129 """Decorator to ensure consistent input/output types for Chebfun __call__ method.
131 This decorator ensures that when a Chebfun object is called with a float input,
132 it returns a float output, and when called with an array input, it returns an
133 array output. It handles various input formats including scalars, numpy arrays,
134 and nested arrays.
136 Args:
137 f (callable): The __call__ method to be wrapped.
139 Returns:
140 callable: A wrapped version of the method that ensures type consistency.
141 """
143 @wraps(f)
144 def thewrapper(self, *args, **kwargs):
145 x = args[0]
146 xx = np.array([x]) if np.isscalar(x) else np.array(x)
147 # discern between the array(0.1) and array([0.1]) cases
148 if xx.size == 1:
149 if xx.ndim == 0:
150 xx = np.array([xx])
151 args = list(args)
152 args[0] = xx
153 out = f(self, *args, **kwargs)
154 return out[0] if np.isscalar(x) else out
156 return thewrapper
159def cast_arg_to_chebfun(f: callable) -> callable:
160 """Decorator to cast the first argument to a chebfun object if needed.
162 This decorator attempts to convert the first argument to a chebfun object
163 if it is not already one. Currently, only numeric types can be cast to chebfun.
165 Args:
166 f (callable): The method to be wrapped.
168 Returns:
169 callable: A wrapped version of the method that ensures the first argument
170 is a chebfun object.
171 """
173 @wraps(f)
174 def wrapper(self, *args, **kwargs):
175 other = args[0]
176 if not isinstance(other, self.__class__):
177 fun = self.initconst(args[0], self.support)
178 args = list(args)
179 args[0] = fun
180 return f(self, *args, **kwargs)
182 return wrapper
185def cast_other(f: callable) -> callable:
186 """Decorator to cast the first argument to the same type as self.
188 This generic decorator is applied to binary operator methods to ensure that
189 the first positional argument (typically 'other') is cast to the same type
190 as the object on which the method is called.
192 Args:
193 f (callable): The binary operator method to be wrapped.
195 Returns:
196 callable: A wrapped version of the method that ensures type consistency
197 between self and the first argument.
198 """
200 @wraps(f)
201 def wrapper(self, *args, **kwargs):
202 cls = self.__class__
203 other = args[0]
204 if not isinstance(other, cls):
205 args = list(args)
206 args[0] = cls(other)
207 return f(self, *args, **kwargs)
209 return wrapper