Coverage for src / tracekit / __main__.py: 100%

1"""TraceKit command-line interface. 

2 

3This module provides the main entry point for TraceKit CLI operations, 

4including sample data download. 

5 

6 

7Example: 

8 python -m tracekit download_samples 

9 python -m tracekit --help 

10""" 

11 

12from __future__ import annotations 

13 

14import argparse 

15import hashlib 

16import sys 

17from pathlib import Path 

18from typing import Any 

19 

20 

21def get_samples_dir() -> Path: 

22 """Get the samples directory path (~/.tracekit/samples/). 

23 

24 Returns: 

25 Path to the samples directory. 

26 """ 

27 return Path.home() / ".tracekit" / "samples" 

28 

29 

30def get_sample_files() -> dict[str, dict[str, Any]]: 

31 """Get the list of sample files to download. 

32 

33 Returns: 

34 Dictionary mapping filename to file metadata. 

35 """ 

36 # Sample files configuration 

37 # In production, these would be hosted on a public CDN or GitHub releases 

38 return { 

39 "sine_1khz.csv": { 

40 "description": "1 kHz sine wave, 100 kS/s, CSV format", 

41 "format": "csv", 

42 "size": 1024 * 50, # ~50 KB 

43 "checksum": None, # Would be populated with actual checksum 

44 "url": "https://raw.githubusercontent.com/tracekit/sample-data/main/sine_1khz.csv", 

45 }, 

46 "square_wave.csv": { 

47 "description": "10 kHz square wave with ringing, CSV format", 

48 "format": "csv", 

49 "size": 1024 * 100, # ~100 KB 

50 "checksum": None, 

51 "url": "https://raw.githubusercontent.com/tracekit/sample-data/main/square_wave.csv", 

52 }, 

53 "uart_9600.bin": { 

54 "description": "UART signal at 9600 baud, binary format", 

55 "format": "binary", 

56 "size": 1024 * 20, # ~20 KB 

57 "checksum": None, 

58 "url": "https://raw.githubusercontent.com/tracekit/sample-data/main/uart_9600.bin", 

59 }, 

60 "i2c_capture.bin": { 

61 "description": "I2C bus capture with multiple devices", 

62 "format": "binary", 

63 "size": 1024 * 50, # ~50 KB 

64 "checksum": None, 

65 "url": "https://raw.githubusercontent.com/tracekit/sample-data/main/i2c_capture.bin", 

66 }, 

67 "spi_flash.bin": { 

68 "description": "SPI flash read operation", 

69 "format": "binary", 

70 "size": 1024 * 30, # ~30 KB 

71 "checksum": None, 

72 "url": "https://raw.githubusercontent.com/tracekit/sample-data/main/spi_flash.bin", 

73 }, 

74 "noisy_signal.csv": { 

75 "description": "Noisy analog signal for filtering examples", 

76 "format": "csv", 

77 "size": 1024 * 80, # ~80 KB 

78 "checksum": None, 

79 "url": "https://raw.githubusercontent.com/tracekit/sample-data/main/noisy_signal.csv", 

80 }, 

81 "eye_diagram.npz": { 

82 "description": "High-speed serial data for eye diagram", 

83 "format": "npz", 

84 "size": 1024 * 200, # ~200 KB 

85 "checksum": None, 

86 "url": "https://raw.githubusercontent.com/tracekit/sample-data/main/eye_diagram.npz", 

87 }, 

88 } 

89 

90 

91def download_file(url: str, dest: Path, checksum: str | None = None) -> bool: 

92 """Download a file from URL to destination. 

93 

94 Args: 

95 url: URL to download from. 

96 dest: Destination file path. 

97 checksum: Optional SHA256 checksum to verify. 

98 

99 Returns: 

100 True if download successful, False otherwise. 

101 """ 

102 try: 

103 import ssl 

104 import urllib.request 

105 

106 # Create SSL context that works in most environments 

107 context = ssl.create_default_context() 

108 

109 print(f" Downloading: {url}") 

110 

111 with urllib.request.urlopen(url, context=context, timeout=30) as response: 

112 data = response.read() 

113 

114 # Verify checksum if provided 

115 if checksum: 

116 computed = hashlib.sha256(data).hexdigest() 

117 if computed != checksum: 

118 print(f" ERROR: Checksum mismatch for {dest.name}") 

119 print(f" Expected: {checksum}") 

120 print(f" Got: {computed}") 

121 return False 

122 

123 # Write to destination 

124 dest.parent.mkdir(parents=True, exist_ok=True) 

125 dest.write_bytes(data) 

126 

127 print(f" Saved: {dest}") 

128 return True 

129 

130 except Exception as e: 

131 print(f" ERROR: Failed to download {url}: {e}") 

132 return False 

133 

134 

135def generate_sample_file(filename: str, dest: Path) -> bool: 

136 """Generate a sample file locally when download is not available. 

137 

138 This is a fallback when the remote repository is not available. 

139 

140 Args: 

141 filename: Name of the sample file to generate. 

142 dest: Destination file path. 

143 

144 Returns: 

145 True if generation successful, False otherwise. 

146 """ 

147 try: 

148 import numpy as np 

149 

150 dest.parent.mkdir(parents=True, exist_ok=True) 

151 

152 if filename == "sine_1khz.csv": 

153 # Generate 1 kHz sine wave at 100 kS/s, 1000 samples 

154 sample_rate = 100_000 

155 duration = 0.01 # 10 ms 

156 t = np.arange(0, duration, 1 / sample_rate) 

157 signal = 0.5 * np.sin(2 * np.pi * 1000 * t) 

158 np.savetxt( 

159 dest, 

160 np.column_stack([t, signal]), 

161 delimiter=",", 

162 header="time,voltage", 

163 comments="", 

164 ) 

165 return True 

166 

167 elif filename == "square_wave.csv": 

168 # Generate 10 kHz square wave with some ringing 

169 sample_rate = 1_000_000 

170 duration = 0.001 # 1 ms 

171 t = np.arange(0, duration, 1 / sample_rate) 

172 signal = 0.5 * np.sign(np.sin(2 * np.pi * 10000 * t)) 

173 # Add some ringing/noise 

174 signal += 0.05 * np.random.randn(len(signal)) 

175 np.savetxt( 

176 dest, 

177 np.column_stack([t, signal]), 

178 delimiter=",", 

179 header="time,voltage", 

180 comments="", 

181 ) 

182 return True 

183 

184 elif filename == "noisy_signal.csv": 

185 # Generate noisy sine wave 

186 sample_rate = 10_000 

187 duration = 0.1 # 100 ms 

188 t = np.arange(0, duration, 1 / sample_rate) 

189 signal = 0.5 * np.sin(2 * np.pi * 100 * t) 

190 signal += 0.1 * np.random.randn(len(signal)) 

191 np.savetxt( 

192 dest, 

193 np.column_stack([t, signal]), 

194 delimiter=",", 

195 header="time,voltage", 

196 comments="", 

197 ) 

198 return True 

199 

200 elif filename.endswith(".bin"): 

201 # Generate placeholder binary data 

202 data = np.random.randint(0, 256, 1000, dtype=np.uint8) 

203 data.tofile(dest) 

204 return True 

205 

206 elif filename.endswith(".npz"): 

207 # Generate high-speed signal for eye diagram 

208 sample_rate = 10_000_000 # 10 MS/s 

209 samples_per_ui = 100 

210 num_ui = 100 

211 t = np.arange(samples_per_ui * num_ui) / sample_rate 

212 # Generate random bit pattern 

213 bits = np.random.randint(0, 2, num_ui) 

214 signal = np.repeat(bits.astype(float), samples_per_ui) 

215 # Add some jitter and noise 

216 signal += 0.1 * np.random.randn(len(signal)) 

217 np.savez(dest, time=t, signal=signal, sample_rate=sample_rate) 

218 return True 

219 

220 else: 

221 print(f" WARNING: Unknown file type: {filename}") 

222 return False 

223 

224 except Exception as e: 

225 print(f" ERROR: Failed to generate {filename}: {e}") 

226 return False 

227 

228 

229def download_samples(force: bool = False, generate: bool = True) -> int: 

230 """Download sample waveform files for testing and tutorials. 

231 

232 Args: 

233 force: Force re-download even if files exist. 

234 generate: Generate files locally if download fails. 

235 

236 Returns: 

237 Exit code (0 for success, 1 for failure). 

238 """ 

239 samples_dir = get_samples_dir() 

240 sample_files = get_sample_files() 

241 

242 print("TraceKit Sample Data Download") 

243 print("==============================") 

244 print(f"Destination: {samples_dir}") 

245 print() 

246 

247 success_count = 0 

248 fail_count = 0 

249 

250 for filename, info in sample_files.items(): 

251 dest = samples_dir / filename 

252 

253 if dest.exists() and not force: 

254 print(f"[SKIP] {filename} (already exists)") 

255 success_count += 1 

256 continue 

257 

258 print(f"[DOWNLOAD] {filename}") 

259 print(f" Description: {info['description']}") 

260 

261 # Try to download first 

262 url = info.get("url") 

263 checksum = info.get("checksum") 

264 

265 if url and download_file(url, dest, checksum): 

266 success_count += 1 

267 continue 

268 

269 # Fall back to local generation 

270 if generate: 

271 print(" Falling back to local generation...") 

272 if generate_sample_file(filename, dest): 

273 print(f" Generated: {dest}") 

274 success_count += 1 

275 continue 

276 

277 fail_count += 1 

278 print(f" FAILED: {filename}") 

279 

280 print() 

281 print(f"Summary: {success_count} succeeded, {fail_count} failed") 

282 

283 if fail_count > 0: 

284 print() 

285 print("Some downloads failed. Sample files are optional and used for") 

286 print("tutorials and testing. You can proceed without them.") 

287 return 1 

288 

289 print() 

290 print("Sample files downloaded successfully!") 

291 print() 

292 print("Example usage:") 

293 print(" >>> import tracekit as tk") 

294 print(f" >>> trace = tk.load('{samples_dir / 'sine_1khz.csv'}')") 

295 print(" >>> trace.plot()") 

296 

297 return 0 

298 

299 

300def list_samples() -> int: 

301 """List available sample files. 

302 

303 Returns: 

304 Exit code. 

305 """ 

306 samples_dir = get_samples_dir() 

307 sample_files = get_sample_files() 

308 

309 print("Available sample files:") 

310 print() 

311 

312 for filename, info in sample_files.items(): 

313 dest = samples_dir / filename 

314 status = "[EXISTS]" if dest.exists() else "[NOT DOWNLOADED]" 

315 print(f" {status} {filename}") 

316 print(f" {info['description']}") 

317 print() 

318 

319 return 0 

320 

321 

322def main() -> int: 

323 """Main entry point for TraceKit CLI. 

324 

325 Returns: 

326 Exit code. 

327 """ 

328 parser = argparse.ArgumentParser( 

329 prog="tracekit", 

330 description="TraceKit signal analysis toolkit", 

331 ) 

332 

333 subparsers = parser.add_subparsers(dest="command", help="Available commands") 

334 

335 # download_samples command 

336 download_parser = subparsers.add_parser( 

337 "download_samples", 

338 aliases=["download"], 

339 help="Download sample waveform files", 

340 ) 

341 download_parser.add_argument( 

342 "--force", 

343 "-f", 

344 action="store_true", 

345 help="Force re-download even if files exist", 

346 ) 

347 download_parser.add_argument( 

348 "--no-generate", 

349 action="store_true", 

350 help="Do not generate files locally if download fails", 

351 ) 

352 

353 # list_samples command 

354 subparsers.add_parser( 

355 "list_samples", 

356 aliases=["list"], 

357 help="List available sample files", 

358 ) 

359 

360 # version command 

361 subparsers.add_parser( 

362 "version", 

363 help="Show version information", 

364 ) 

365 

366 args = parser.parse_args() 

367 

368 if args.command in ("download_samples", "download"): 

369 return download_samples( 

370 force=args.force, 

371 generate=not args.no_generate, 

372 ) 

373 

374 elif args.command in ("list_samples", "list"): 

375 return list_samples() 

376 

377 elif args.command == "version": 

378 try: 

379 from tracekit import __version__ 

380 

381 print(f"TraceKit version {__version__}") 

382 except ImportError: 

383 print("TraceKit version unknown") 

384 return 0 

385 

386 else: 

387 parser.print_help() 

388 return 0 

389 

390 

391if __name__ == "__main__": 

392 sys.exit(main())