Coverage for src/chebpy/decorators.py: 100%
76 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-22 21:50 +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 collections.abc import Callable
11from functools import wraps
12from typing import Any
14import numpy as np
17def cache(f: Callable[..., Any]) -> Callable[..., Any]:
18 """Object method output caching mechanism.
20 This decorator caches the output of zero-argument methods to speed up repeated
21 execution of relatively expensive operations such as .roots(). Cached computations
22 are stored in a dictionary called _cache which is bound to self using keys
23 corresponding to the method name.
25 Args:
26 f (callable): The method to be cached. Must be a zero-argument method.
28 Returns:
29 callable: A wrapped version of the method that implements caching.
31 Note:
32 Can be used in principle on arbitrary objects.
33 """
35 # TODO: look into replacing this with one of the functools cache decorators
36 @wraps(f)
37 def wrapper(self: Any) -> Any:
38 """Return the cached method result, computing and storing it on first call."""
39 try:
40 # f has been executed previously
41 out = self._cache[f.__name__] # ty: ignore[unresolved-attribute]
42 except AttributeError:
43 # f has not been executed previously and self._cache does not exist
44 self._cache = {}
45 out = self._cache[f.__name__] = f(self) # ty: ignore[unresolved-attribute]
46 except KeyError: # pragma: no cover
47 # f has not been executed previously, but self._cache exists
48 out = self._cache[f.__name__] = f(self) # ty: ignore[unresolved-attribute]
49 return out
51 return wrapper
54def self_empty(resultif: Any = None) -> Callable[..., Any]:
55 """Factory method to produce a decorator for handling empty objects.
57 This factory creates a decorator that checks whether the object whose method
58 is being wrapped is empty. If the object is empty, it returns either the supplied
59 resultif value or a copy of the object. Otherwise, it executes the wrapped method.
61 Args:
62 resultif: Value to return when the object is empty. If None, returns a copy
63 of the object instead.
65 Returns:
66 callable: A decorator function that implements the empty-checking logic.
68 Note:
69 This decorator is primarily used in chebtech.py.
70 """
72 # TODO: add unit test for this
73 def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
74 """Wrap *f* with the empty-object short-circuit logic."""
76 @wraps(f)
77 def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
78 """Return the empty-case result if *self* is empty, else call *f*."""
79 if self.isempty:
80 if resultif is not None:
81 return resultif
82 else:
83 return self.copy()
84 else:
85 return f(self, *args, **kwargs)
87 return wrapper
89 return decorator
92def preandpostprocess(f: Callable[..., Any]) -> Callable[..., Any]:
93 """Decorator for pre- and post-processing tasks common to bary and clenshaw.
95 This decorator handles several edge cases for functions like bary and clenshaw:
96 - Empty arrays in input arguments
97 - Constant functions
98 - NaN values in coefficients
99 - Scalar vs. array inputs
101 Args:
102 f (callable): The function to be wrapped.
104 Returns:
105 callable: A wrapped version of the function with pre- and post-processing.
106 """
108 @wraps(f)
109 def thewrapper(*args: Any, **kwargs: Any) -> Any:
110 """Handle empty/constant/NaN/scalar edge cases around *f*."""
111 xx, akfk = args[:2]
112 # are any of the first two arguments empty arrays?
113 if (np.asarray(xx).size == 0) | (np.asarray(akfk).size == 0):
114 return np.array([])
115 # is the function constant?
116 elif akfk.size == 1:
117 if np.isscalar(xx):
118 return akfk[0]
119 else:
120 return akfk * np.ones(xx.size)
121 # are there any NaNs in the second argument?
122 elif np.any(np.isnan(akfk)):
123 return np.nan * np.ones(xx.size)
124 # convert first argument to an array if it is a scalar and then
125 # return the first (and only) element of the result if so
126 else:
127 args_list = list(args)
128 args_list[0] = np.array([xx]) if np.isscalar(xx) else args_list[0]
129 out = f(*args_list, **kwargs)
130 return out[0] if np.isscalar(xx) else out
132 return thewrapper
135def float_argument(f: Callable[..., Any]) -> Callable[..., Any]:
136 """Decorator to ensure consistent input/output types for Chebfun __call__ method.
138 This decorator ensures that when a Chebfun object is called with a float input,
139 it returns a float output, and when called with an array input, it returns an
140 array output. It handles various input formats including scalars, numpy arrays,
141 and nested arrays.
143 Args:
144 f (callable): The __call__ method to be wrapped.
146 Returns:
147 callable: A wrapped version of the method that ensures type consistency.
148 """
150 @wraps(f)
151 def thewrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
152 """Coerce the first argument to an array and match scalar/array output to it."""
153 x = args[0]
154 xx = np.array([x]) if np.isscalar(x) else np.array(x)
155 # discern between the array(0.1) and array([0.1]) cases
156 if xx.size == 1 and xx.ndim == 0:
157 xx = np.array([xx])
158 args_list = list(args)
159 args_list[0] = xx
160 out = f(self, *args_list, **kwargs)
161 return out[0] if np.isscalar(x) else out
163 return thewrapper
166def cast_arg_to_chebfun(f: Callable[..., Any]) -> Callable[..., Any]:
167 """Decorator to cast the first argument to a chebfun object if needed.
169 This decorator attempts to convert the first argument to a chebfun object
170 if it is not already one. Currently, only numeric types can be cast to chebfun.
172 Args:
173 f (callable): The method to be wrapped.
175 Returns:
176 callable: A wrapped version of the method that ensures the first argument
177 is a chebfun object.
178 """
180 @wraps(f)
181 def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
182 """Cast the first argument to a chebfun (if needed) before calling *f*."""
183 other = args[0]
184 if not isinstance(other, self.__class__):
185 fun = self.initconst(args[0], self.support)
186 args_list = list(args)
187 args_list[0] = fun
188 return f(self, *args_list, **kwargs)
189 return f(self, *args, **kwargs)
191 return wrapper
194def cast_other(f: Callable[..., Any]) -> Callable[..., Any]:
195 """Decorator to cast the first argument to the same type as self.
197 This generic decorator is applied to binary operator methods to ensure that
198 the first positional argument (typically 'other') is cast to the same type
199 as the object on which the method is called.
201 Args:
202 f (callable): The binary operator method to be wrapped.
204 Returns:
205 callable: A wrapped version of the method that ensures type consistency
206 between self and the first argument.
207 """
209 @wraps(f)
210 def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
211 """Cast the first argument to ``type(self)`` (if needed) before calling *f*."""
212 cls = self.__class__
213 other = args[0]
214 if not isinstance(other, cls):
215 args_list = list(args)
216 args_list[0] = cls(other)
217 return f(self, *args_list, **kwargs)
218 return f(self, *args, **kwargs)
220 return wrapper