Coverage for src / tracekit / jupyter / magic.py: 15%

1"""IPython magic commands for TraceKit. 

2 

3This module provides IPython/Jupyter magic commands for convenient 

4trace analysis in notebooks. 

5 

6 - %tracekit load <file> - Load a trace file 

7 - %tracekit measure - Run measurements on current trace 

8 - %%analyze - Multi-line analysis cell 

9 - Auto-display of results with rich HTML 

10 

11Example: 

12 In [1]: %load_ext tracekit 

13 

14 In [2]: %tracekit load capture.wfm 

15 Loaded: WaveformTrace with 10000 samples @ 1 GSa/s 

16 

17 In [3]: %tracekit measure rise_time fall_time 

18 rise_time: 2.5 ns 

19 fall_time: 2.8 ns 

20 

21 In [4]: %%analyze 

22 ...: trace = load("capture.wfm") 

23 ...: print(f"THD: {thd(trace):.2f} dB") 

24 

25References: 

26 - IPython Magic Commands documentation 

27 - Jupyter display architecture 

28""" 

29 

30from __future__ import annotations 

31 

32from typing import TYPE_CHECKING, Any 

33 

34if TYPE_CHECKING: 

35 from IPython.core.interactiveshell import InteractiveShell 

36 

37 

38# Store current trace for magic commands 

39_current_trace: Any = None 

40_current_file: str | None = None 

41 

42 

43def get_current_trace() -> Any: 

44 """Get the currently loaded trace from magic commands.""" 

45 return _current_trace 

46 

47 

48def set_current_trace(trace: Any, filename: str | None = None) -> None: 

49 """Set the current trace for magic commands.""" 

50 global _current_trace, _current_file 

51 _current_trace = trace 

52 _current_file = filename 

53 

54 

55try: 

56 from IPython.core.magic import ( 

57 Magics, 

58 cell_magic, 

59 line_magic, 

60 magics_class, 

61 ) 

62 from IPython.display import HTML, display # noqa: F401 

63 

64 IPYTHON_AVAILABLE = True 

65except ImportError: 

66 IPYTHON_AVAILABLE = False 

67 

68 class Magics: # type: ignore[no-redef] 

69 """Fallback Magics class when IPython not available.""" 

70 

71 pass 

72 

73 def magics_class(cls: type[Any]) -> type[Any]: # type: ignore[no-redef,misc] 

74 """Dummy decorator when IPython not available.""" 

75 return cls 

76 

77 def line_magic(func): # type: ignore[no-untyped-def] 

78 """Dummy decorator.""" 

79 return func 

80 

81 def cell_magic(func): # type: ignore[no-untyped-def] 

82 """Dummy decorator.""" 

83 return func 

84 

85 

86@magics_class 

87class TracekitMagics(Magics): # type: ignore[misc] 

88 """IPython magics for TraceKit analysis. 

89 

90 Provides convenient shortcuts for loading traces, running measurements, 

91 and displaying results in Jupyter notebooks. 

92 """ 

93 

94 @line_magic # type: ignore[misc, untyped-decorator] 

95 def tracekit(self, line: str) -> Any: 

96 """TraceKit line magic for quick operations. 

97 

98 Usage: 

99 %tracekit load <filename> - Load a trace file 

100 %tracekit measure [names...] - Run measurements 

101 %tracekit info - Show current trace info 

102 %tracekit formats - List supported formats 

103 %tracekit help - Show help 

104 

105 Args: 

106 line: Magic command arguments 

107 

108 Returns: 

109 Command result or None 

110 """ 

111 import tracekit as tk 

112 

113 parts = line.strip().split() 

114 if not parts: 

115 return self._show_help() 

116 

117 cmd = parts[0].lower() 

118 

119 if cmd == "load": 

120 if len(parts) < 2: 

121 print("Usage: %tracekit load <filename>") 

122 return None 

123 filename = " ".join(parts[1:]) 

124 return self._load_trace(filename) 

125 

126 elif cmd == "measure": 

127 measurements = parts[1:] if len(parts) > 1 else None 

128 return self._run_measurements(measurements) 

129 

130 elif cmd == "info": 

131 return self._show_trace_info() 

132 

133 elif cmd == "formats": 

134 formats = tk.get_supported_formats() 

135 print("Supported formats:") 

136 for fmt in formats: 

137 print(f" - {fmt}") 

138 return formats 

139 

140 elif cmd == "help": 

141 return self._show_help() 

142 

143 else: 

144 print(f"Unknown command: {cmd}") 

145 return self._show_help() 

146 

147 def _load_trace(self, filename: str) -> Any: 

148 """Load a trace file and store as current.""" 

149 import tracekit as tk 

150 

151 try: 

152 trace = tk.load(filename) 

153 set_current_trace(trace, filename) 

154 

155 # Display summary 

156 info = { 

157 "file": filename, 

158 "type": type(trace).__name__, 

159 "samples": len(trace.data) if hasattr(trace, "data") else "N/A", 

160 "sample_rate": f"{trace.metadata.sample_rate / 1e9:.3f} GSa/s" 

161 if hasattr(trace, "metadata") 

162 else "N/A", 

163 } 

164 

165 print(f"Loaded: {info['type']} with {info['samples']} samples @ {info['sample_rate']}") 

166 return trace 

167 

168 except Exception as e: 

169 print(f"Error loading {filename}: {e}") 

170 return None 

171 

172 def _run_measurements(self, measurement_names: list[str] | None) -> dict[str, Any]: 

173 """Run measurements on current trace.""" 

174 import tracekit as tk 

175 

176 trace = get_current_trace() 

177 if trace is None: 

178 print("No trace loaded. Use: %tracekit load <filename>") 

179 return {} 

180 

181 if measurement_names: 

182 # Run specific measurements 

183 results = {} 

184 for name in measurement_names: 

185 if hasattr(tk, name): 

186 try: 

187 func = getattr(tk, name) 

188 results[name] = func(trace) 

189 except Exception as e: 

190 results[name] = f"Error: {e}" 

191 else: 

192 results[name] = "Unknown measurement" 

193 else: 

194 # Run all measurements 

195 try: 

196 results = tk.measure(trace) 

197 except Exception as e: 

198 print(f"Error running measurements: {e}") 

199 return {} 

200 

201 # Display results 

202 self._display_measurements(results) 

203 return results 

204 

205 def _display_measurements(self, results: dict[str, Any]) -> None: 

206 """Display measurement results with formatting.""" 

207 if IPYTHON_AVAILABLE: 

208 from tracekit.jupyter.display import display_measurements 

209 

210 display_measurements(results) 

211 else: 

212 for name, value in results.items(): 

213 if isinstance(value, float): 

214 print(f"{name}: {value:.6g}") 

215 else: 

216 print(f"{name}: {value}") 

217 

218 def _show_trace_info(self) -> dict[str, Any] | None: 

219 """Show information about current trace.""" 

220 trace = get_current_trace() 

221 if trace is None: 

222 print("No trace loaded. Use: %tracekit load <filename>") 

223 return None 

224 

225 info = { 

226 "file": _current_file, 

227 "type": type(trace).__name__, 

228 } 

229 

230 if hasattr(trace, "data"): 

231 info["samples"] = len(trace.data) # type: ignore[assignment] 

232 

233 if hasattr(trace, "metadata"): 

234 meta = trace.metadata 

235 if hasattr(meta, "sample_rate"): 

236 info["sample_rate"] = meta.sample_rate 

237 if hasattr(meta, "channel_name"): 

238 info["channel"] = meta.channel_name 

239 

240 for key, value in info.items(): 

241 print(f"{key}: {value}") 

242 

243 return info 

244 

245 def _show_help(self) -> None: 

246 """Show magic command help.""" 

247 help_text = """ 

248TraceKit Magic Commands 

249======================= 

250 

251%tracekit load <filename> Load a trace file 

252%tracekit measure [names...] Run measurements on current trace 

253%tracekit info Show current trace info 

254%tracekit formats List supported file formats 

255%tracekit help Show this help 

256 

257%%analyze Multi-line analysis cell magic 

258 

259Available measurements: rise_time, fall_time, frequency, period, 

260 amplitude, rms, overshoot, undershoot, thd, snr, sinad 

261 

262Example: 

263 %tracekit load capture.wfm 

264 %tracekit measure rise_time fall_time 

265""" 

266 print(help_text) 

267 

268 @cell_magic # type: ignore[misc, untyped-decorator] 

269 def analyze(self, line: str, cell: str) -> Any: 

270 """Cell magic for multi-line analysis. 

271 

272 Args: 

273 line: Magic command line arguments (unused). 

274 cell: Multi-line cell content to execute. 

275 

276 Returns: 

277 Result from cell execution (if 'result' variable defined). 

278 

279 Usage: 

280 %%analyze 

281 trace = load("capture.wfm") 

282 result = measure(trace) 

283 print(f"Rise time: {result['rise_time']}") 

284 

285 All tracekit functions are auto-imported in the cell namespace. 

286 """ 

287 import tracekit as tk 

288 

289 # Build execution namespace with tracekit imports 

290 namespace = { 

291 "tk": tk, 

292 "load": tk.load, 

293 "measure": tk.measure, 

294 "fft": tk.fft, 

295 "psd": tk.psd, 

296 "thd": tk.thd, 

297 "snr": tk.snr, 

298 "rise_time": tk.rise_time, 

299 "fall_time": tk.fall_time, 

300 "frequency": tk.frequency, 

301 "amplitude": tk.amplitude, 

302 "low_pass": tk.low_pass, 

303 "high_pass": tk.high_pass, 

304 } 

305 

306 # Add current trace if available 

307 trace = get_current_trace() 

308 if trace is not None: 

309 namespace["trace"] = trace 

310 

311 # Execute cell 

312 exec(cell, namespace) 

313 

314 # Return any result variable if defined 

315 return namespace.get("result") 

316 

317 

318def load_ipython_extension(ipython: InteractiveShell) -> None: 

319 """Load TraceKit IPython extension. 

320 

321 Called when user runs: %load_ext tracekit 

322 

323 Args: 

324 ipython: The IPython shell instance 

325 """ 

326 ipython.register_magics(TracekitMagics) 

327 print("TraceKit magics loaded. Type '%tracekit help' for usage.") 

328 

329 

330def unload_ipython_extension(ipython: InteractiveShell) -> None: 

331 """Unload TraceKit IPython extension. 

332 

333 Args: 

334 ipython: The IPython shell instance 

335 """ 

336 # IPython handles magic cleanup automatically 

337 pass