Coverage for src / tracekit / extensibility / registry.py: 33%

1"""Algorithm registry for custom algorithm injection. 

2 

3This module implements a registry pattern that allows users to register 

4custom algorithms and implementations at extension points throughout TraceKit. 

5""" 

6 

7from __future__ import annotations 

8 

9import inspect 

10from typing import TYPE_CHECKING, Any 

11 

12if TYPE_CHECKING: 

13 from collections.abc import Callable 

14 

15 

16class AlgorithmRegistry: 

17 """Singleton registry for custom algorithm implementations. 

18 

19 Allows users to register custom algorithms for various categories 

20 (edge detection, peak finding, window functions, etc.) that can be 

21 used throughout TraceKit. 

22 

23 The registry validates algorithm signatures on registration and provides 

24 lookup by category and name. 

25 

26 Example: 

27 >>> import tracekit as tk 

28 >>> def my_edge_detector(data, threshold=0.5): 

29 ... '''Custom Schmitt trigger edge detector''' 

30 ... edges = [] 

31 ... state = data[0] > threshold 

32 ... for i, val in enumerate(data): 

33 ... new_state = val > threshold 

34 ... if new_state != state: 

35 ... edges.append(i) 

36 ... state = new_state 

37 ... return edges 

38 >>> # Register in algorithm registry 

39 >>> tk.register_algorithm('my_schmitt', my_edge_detector, category='edge_detector') 

40 >>> # Use custom algorithm 

41 >>> edges = tk.find_edges(trace, method='my_schmitt', threshold=0.7) 

42 

43 Advanced Example: 

44 >>> # Register custom window function 

45 >>> import numpy as np 

46 >>> def custom_window(n, alpha=0.5): 

47 ... x = np.linspace(0, 1, n) 

48 ... return 0.5 * (1 + np.cos(2 * np.pi * alpha * (x - 0.5))) 

49 >>> tk.register_algorithm('custom_tukey', custom_window, category='window_func') 

50 >>> # Use in FFT 

51 >>> result = tk.fft(trace, nfft=8192, window='custom_tukey', alpha=0.3) 

52 >>> # List available algorithms 

53 >>> available = tk.get_algorithms('window_func') 

54 

55 References: 

56 API-006: Algorithm Override Hooks 

57 pytest plugin system 

58 https://docs.pytest.org/en/stable/how-to/writing_plugins.html 

59 """ 

60 

61 _instance: AlgorithmRegistry | None = None 

62 _registries: dict[str, dict[str, Callable[..., Any]]] 

63 

64 def __new__(cls) -> AlgorithmRegistry: 

65 """Ensure singleton instance.""" 

66 if cls._instance is None: 

67 cls._instance = super().__new__(cls) 

68 cls._instance._registries = {} 

69 return cls._instance 

70 

71 def register( 

72 self, 

73 name: str, 

74 func: Callable[..., Any], 

75 category: str, 

76 validate: bool = True, 

77 ) -> None: 

78 """Register a custom algorithm. 

79 

80 Args: 

81 name: Unique name for the algorithm within its category. 

82 func: Callable implementing the algorithm. 

83 category: Category of algorithm (e.g., 'edge_detector', 'peak_finder'). 

84 validate: Whether to validate function signature. Default True. 

85 

86 Raises: 

87 ValueError: If name already exists in category. 

88 TypeError: If func is not callable or signature is invalid. 

89 

90 Example: 

91 >>> def my_algorithm(data, param1=1.0, param2=2.0): 

92 ... return data * param1 + param2 

93 >>> registry = AlgorithmRegistry() 

94 >>> registry.register('my_algo', my_algorithm, 'preprocessor') 

95 """ 

96 if not callable(func): 

97 raise TypeError(f"Algorithm must be callable, got {type(func).__name__}") 

98 

99 # Initialize category if needed 

100 if category not in self._registries: 

101 self._registries[category] = {} 

102 

103 # Check for duplicates 

104 if name in self._registries[category]: 

105 raise ValueError(f"Algorithm '{name}' already registered in category '{category}'") 

106 

107 # Validate signature if requested 

108 if validate: 

109 self._validate_signature(func, category) 

110 

111 # Register algorithm 

112 self._registries[category][name] = func 

113 

114 def get(self, category: str, name: str) -> Callable[..., Any]: 

115 """Get algorithm by category and name. 

116 

117 Args: 

118 category: Algorithm category. 

119 name: Algorithm name. 

120 

121 Returns: 

122 The registered algorithm function. 

123 

124 Raises: 

125 KeyError: If category or name not found. 

126 

127 Example: 

128 >>> registry = AlgorithmRegistry() 

129 >>> edge_detector = registry.get('edge_detector', 'my_schmitt') 

130 >>> edges = edge_detector(data, threshold=0.5) 

131 """ 

132 if category not in self._registries: 

133 raise KeyError( 

134 f"Category '{category}' not found. Available: {list(self._registries.keys())}" 

135 ) 

136 

137 if name not in self._registries[category]: 

138 raise KeyError( 

139 f"Algorithm '{name}' not found in category '{category}'. " 

140 f"Available: {list(self._registries[category].keys())}" 

141 ) 

142 

143 return self._registries[category][name] 

144 

145 def list_categories(self) -> list[str]: 

146 """List all registered algorithm categories. 

147 

148 Returns: 

149 List of category names. 

150 

151 Example: 

152 >>> registry = AlgorithmRegistry() 

153 >>> categories = registry.list_categories() 

154 >>> print(categories) 

155 ['edge_detector', 'peak_finder', 'window_func'] 

156 """ 

157 return list(self._registries.keys()) 

158 

159 def list_algorithms(self, category: str) -> list[str]: 

160 """List all algorithms in a category. 

161 

162 Args: 

163 category: Algorithm category. 

164 

165 Returns: 

166 List of algorithm names in that category. 

167 

168 Raises: 

169 KeyError: If category not found. 

170 

171 Example: 

172 >>> registry = AlgorithmRegistry() 

173 >>> algorithms = registry.list_algorithms('edge_detector') 

174 >>> print(algorithms) 

175 ['threshold', 'hysteresis', 'my_schmitt'] 

176 """ 

177 if category not in self._registries: 

178 raise KeyError( 

179 f"Category '{category}' not found. Available: {list(self._registries.keys())}" 

180 ) 

181 

182 return list(self._registries[category].keys()) 

183 

184 def has_algorithm(self, category: str, name: str) -> bool: 

185 """Check if algorithm is registered. 

186 

187 Args: 

188 category: Algorithm category. 

189 name: Algorithm name. 

190 

191 Returns: 

192 True if algorithm is registered. 

193 

194 Example: 

195 >>> if registry.has_algorithm('edge_detector', 'my_schmitt'): 

196 ... detector = registry.get('edge_detector', 'my_schmitt') 

197 """ 

198 return category in self._registries and name in self._registries[category] 

199 

200 def unregister(self, category: str, name: str) -> None: 

201 """Remove algorithm from registry. 

202 

203 Args: 

204 category: Algorithm category. 

205 name: Algorithm name. 

206 

207 Raises: 

208 KeyError: If algorithm not found. 

209 

210 Example: 

211 >>> registry.unregister('edge_detector', 'my_schmitt') 

212 """ 

213 if not self.has_algorithm(category, name): 

214 raise KeyError(f"Algorithm '{name}' not found in category '{category}'") 

215 

216 del self._registries[category][name] 

217 

218 def clear_category(self, category: str) -> None: 

219 """Clear all algorithms in a category. 

220 

221 Args: 

222 category: Category to clear. 

223 

224 Example: 

225 >>> registry.clear_category('edge_detector') 

226 """ 

227 if category in self._registries: 

228 self._registries[category].clear() 

229 

230 def clear_all(self) -> None: 

231 """Clear all registered algorithms. 

232 

233 Example: 

234 >>> registry.clear_all() 

235 """ 

236 self._registries.clear() 

237 

238 def _validate_signature(self, func: Callable[..., Any], category: str) -> None: 

239 """Validate function signature for category. 

240 

241 Args: 

242 func: Function to validate. 

243 category: Category to validate against. 

244 

245 Raises: 

246 TypeError: If signature is invalid for category. 

247 """ 

248 sig = inspect.signature(func) 

249 

250 # Check that function accepts **kwargs for extensibility 

251 has_var_keyword = any( 

252 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() 

253 ) 

254 

255 if not has_var_keyword: 

256 # Check if function has at least one parameter 

257 if len(sig.parameters) == 0: 

258 raise TypeError( 

259 f"Algorithm function must accept at least one parameter " 

260 f"(got {func.__name__} with no parameters)" 

261 ) 

262 

263 def __repr__(self) -> str: 

264 """String representation of registry.""" 

265 total_algos = sum(len(algos) for algos in self._registries.values()) 

266 return f"AlgorithmRegistry(categories={len(self._registries)}, algorithms={total_algos})" 

267 

268 

269# Global registry instance 

270_registry = AlgorithmRegistry() 

271 

272 

273def register_algorithm( 

274 name: str, 

275 func: Callable[..., Any], 

276 category: str, 

277 validate: bool = True, 

278) -> None: 

279 """Register a custom algorithm in the global registry. 

280 

281 Convenience function for registering algorithms without accessing 

282 the registry instance directly. 

283 

284 Args: 

285 name: Unique name for the algorithm. 

286 func: Callable implementing the algorithm. 

287 category: Algorithm category. 

288 validate: Whether to validate signature. Default True. 

289 

290 Example: 

291 >>> import tracekit as tk 

292 >>> def my_edge_detector(data, threshold=0.5): 

293 ... return find_edges_custom(data, threshold) 

294 >>> tk.register_algorithm('my_edges', my_edge_detector, 'edge_detector') 

295 

296 References: 

297 API-006: Algorithm Override Hooks 

298 """ 

299 _registry.register(name, func, category, validate) 

300 

301 

302def get_algorithm(category: str, name: str) -> Callable[..., Any]: 

303 """Get algorithm from global registry. 

304 

305 Args: 

306 category: Algorithm category. 

307 name: Algorithm name. 

308 

309 Returns: 

310 The registered algorithm function. 

311 

312 Example: 

313 >>> edge_detector = tk.get_algorithm('edge_detector', 'my_edges') 

314 >>> edges = edge_detector(data) 

315 

316 References: 

317 API-006: Algorithm Override Hooks 

318 """ 

319 return _registry.get(category, name) 

320 

321 

322def get_algorithms(category: str) -> list[str]: 

323 """List all algorithms in a category from global registry. 

324 

325 Args: 

326 category: Algorithm category. 

327 

328 Returns: 

329 List of algorithm names. 

330 

331 Example: 

332 >>> available = tk.get_algorithms('window_func') 

333 >>> print(f"Available windows: {available}") 

334 

335 References: 

336 API-006: Algorithm Override Hooks 

337 """ 

338 return _registry.list_algorithms(category) 

339 

340 

341__all__ = [ 

342 "AlgorithmRegistry", 

343 "get_algorithm", 

344 "get_algorithms", 

345 "register_algorithm", 

346]