Coverage for src / chebpy / decorators.py: 100%

1"""Decorator functions for the ChebPy package. 

2 

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""" 

9 

10from collections.abc import Callable 

11from functools import wraps 

12from typing import Any 

13 

14import numpy as np 

15 

16 

17def cache(f: Callable[..., Any]) -> Callable[..., Any]: 

18 """Object method output caching mechanism. 

19 

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. 

24 

25 Args: 

26 f (callable): The method to be cached. Must be a zero-argument method. 

27 

28 Returns: 

29 callable: A wrapped version of the method that implements caching. 

30 

31 Note: 

32 Can be used in principle on arbitrary objects. 

33 """ 

34 

35 # TODO: look into replacing this with one of the functools cache decorators 

36 @wraps(f) 

37 def wrapper(self: Any) -> Any: 

38 try: 

39 # f has been executed previously 

40 out = self._cache[f.__name__] # type: ignore[attr-defined] 

41 except AttributeError: 

42 # f has not been executed previously and self._cache does not exist 

43 self._cache = {} 

44 out = self._cache[f.__name__] = f(self) # type: ignore[attr-defined] 

45 except KeyError: # pragma: no cover 

46 # f has not been executed previously, but self._cache exists 

47 out = self._cache[f.__name__] = f(self) # type: ignore[attr-defined] 

48 return out 

49 

50 return wrapper 

51 

52 

53def self_empty(resultif: Any = None) -> Callable[..., Any]: 

54 """Factory method to produce a decorator for handling empty objects. 

55 

56 This factory creates a decorator that checks whether the object whose method 

57 is being wrapped is empty. If the object is empty, it returns either the supplied 

58 resultif value or a copy of the object. Otherwise, it executes the wrapped method. 

59 

60 Args: 

61 resultif: Value to return when the object is empty. If None, returns a copy 

62 of the object instead. 

63 

64 Returns: 

65 callable: A decorator function that implements the empty-checking logic. 

66 

67 Note: 

68 This decorator is primarily used in chebtech.py. 

69 """ 

70 

71 # TODO: add unit test for this 

72 def decorator(f: Callable[..., Any]) -> Callable[..., Any]: 

73 @wraps(f) 

74 def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: 

75 if self.isempty: 

76 if resultif is not None: 

77 return resultif 

78 else: 

79 return self.copy() 

80 else: 

81 return f(self, *args, **kwargs) 

82 

83 return wrapper 

84 

85 return decorator 

86 

87 

88def preandpostprocess(f: Callable[..., Any]) -> Callable[..., Any]: 

89 """Decorator for pre- and post-processing tasks common to bary and clenshaw. 

90 

91 This decorator handles several edge cases for functions like bary and clenshaw: 

92 - Empty arrays in input arguments 

93 - Constant functions 

94 - NaN values in coefficients 

95 - Scalar vs. array inputs 

96 

97 Args: 

98 f (callable): The function to be wrapped. 

99 

100 Returns: 

101 callable: A wrapped version of the function with pre- and post-processing. 

102 """ 

103 

104 @wraps(f) 

105 def thewrapper(*args: Any, **kwargs: Any) -> Any: 

106 xx, akfk = args[:2] 

107 # are any of the first two arguments empty arrays? 

108 if (np.asarray(xx).size == 0) | (np.asarray(akfk).size == 0): 

109 return np.array([]) 

110 # is the function constant? 

111 elif akfk.size == 1: 

112 if np.isscalar(xx): 

113 return akfk[0] 

114 else: 

115 return akfk * np.ones(xx.size) 

116 # are there any NaNs in the second argument? 

117 elif np.any(np.isnan(akfk)): 

118 return np.nan * np.ones(xx.size) 

119 # convert first argument to an array if it is a scalar and then 

120 # return the first (and only) element of the result if so 

121 else: 

122 args_list = list(args) 

123 args_list[0] = np.array([xx]) if np.isscalar(xx) else args_list[0] 

124 out = f(*args_list, **kwargs) 

125 return out[0] if np.isscalar(xx) else out 

126 

127 return thewrapper 

128 

129 

130def float_argument(f: Callable[..., Any]) -> Callable[..., Any]: 

131 """Decorator to ensure consistent input/output types for Chebfun __call__ method. 

132 

133 This decorator ensures that when a Chebfun object is called with a float input, 

134 it returns a float output, and when called with an array input, it returns an 

135 array output. It handles various input formats including scalars, numpy arrays, 

136 and nested arrays. 

137 

138 Args: 

139 f (callable): The __call__ method to be wrapped. 

140 

141 Returns: 

142 callable: A wrapped version of the method that ensures type consistency. 

143 """ 

144 

145 @wraps(f) 

146 def thewrapper(self: Any, *args: Any, **kwargs: Any) -> Any: 

147 x = args[0] 

148 xx = np.array([x]) if np.isscalar(x) else np.array(x) 

149 # discern between the array(0.1) and array([0.1]) cases 

150 if xx.size == 1 and xx.ndim == 0: 

151 xx = np.array([xx]) 

152 args_list = list(args) 

153 args_list[0] = xx 

154 out = f(self, *args_list, **kwargs) 

155 return out[0] if np.isscalar(x) else out 

156 

157 return thewrapper 

158 

159 

160def cast_arg_to_chebfun(f: Callable[..., Any]) -> Callable[..., Any]: 

161 """Decorator to cast the first argument to a chebfun object if needed. 

162 

163 This decorator attempts to convert the first argument to a chebfun object 

164 if it is not already one. Currently, only numeric types can be cast to chebfun. 

165 

166 Args: 

167 f (callable): The method to be wrapped. 

168 

169 Returns: 

170 callable: A wrapped version of the method that ensures the first argument 

171 is a chebfun object. 

172 """ 

173 

174 @wraps(f) 

175 def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: 

176 other = args[0] 

177 if not isinstance(other, self.__class__): 

178 fun = self.initconst(args[0], self.support) 

179 args_list = list(args) 

180 args_list[0] = fun 

181 return f(self, *args_list, **kwargs) 

182 return f(self, *args, **kwargs) 

183 

184 return wrapper 

185 

186 

187def cast_other(f: Callable[..., Any]) -> Callable[..., Any]: 

188 """Decorator to cast the first argument to the same type as self. 

189 

190 This generic decorator is applied to binary operator methods to ensure that 

191 the first positional argument (typically 'other') is cast to the same type 

192 as the object on which the method is called. 

193 

194 Args: 

195 f (callable): The binary operator method to be wrapped. 

196 

197 Returns: 

198 callable: A wrapped version of the method that ensures type consistency 

199 between self and the first argument. 

200 """ 

201 

202 @wraps(f) 

203 def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: 

204 cls = self.__class__ 

205 other = args[0] 

206 if not isinstance(other, cls): 

207 args_list = list(args) 

208 args_list[0] = cls(other) 

209 return f(self, *args_list, **kwargs) 

210 return f(self, *args, **kwargs) 

211 

212 return wrapper