recorder.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183
  1. from __future__ import annotations
  2. from dataclasses import asdict, dataclass
  3. from datetime import datetime
  4. from pathlib import Path
  5. import queue
  6. import threading
  7. import wave
  8. import numpy as np
  9. @dataclass
  10. class RecorderStatus:
  11. recording: bool = False
  12. filename: str = ""
  13. duration_sec: float = 0.0
  14. channels: int = 1
  15. sample_rate: int = 16000
  16. source: str = "mic1"
  17. class WavRecorder:
  18. """Threaded WAV writer to avoid blocking the audio callback."""
  19. def __init__(self, recordings_dir: Path) -> None:
  20. self.recordings_dir = Path(recordings_dir)
  21. self.recordings_dir.mkdir(parents=True, exist_ok=True)
  22. self._lock = threading.Lock()
  23. self._queue: queue.Queue[bytes] = queue.Queue(maxsize=512)
  24. self._stop_event = threading.Event()
  25. self._worker: threading.Thread | None = None
  26. self._wave_file: wave.Wave_write | None = None
  27. self._frames_written = 0
  28. self._status = RecorderStatus()
  29. def get_status(self) -> RecorderStatus:
  30. with self._lock:
  31. return RecorderStatus(**asdict(self._status))
  32. def start(self, source: str, sample_rate: int, channels: int, filename: str | None = None) -> str:
  33. with self._lock:
  34. if self._status.recording:
  35. raise RuntimeError("Recording is already active")
  36. source_clean = source.replace(" ", "_").replace("/", "_")
  37. if filename is None:
  38. timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
  39. filename = f"rec_{timestamp}_{source_clean}.wav"
  40. filepath = self.recordings_dir / filename
  41. self._wave_file = wave.open(str(filepath), "wb")
  42. self._wave_file.setnchannels(channels)
  43. self._wave_file.setsampwidth(2)
  44. self._wave_file.setframerate(sample_rate)
  45. self._stop_event.clear()
  46. self._frames_written = 0
  47. self._clear_queue_locked()
  48. self._status = RecorderStatus(
  49. recording=True,
  50. filename=filename,
  51. duration_sec=0.0,
  52. channels=channels,
  53. sample_rate=sample_rate,
  54. source=source_clean,
  55. )
  56. self._worker = threading.Thread(target=self._writer_loop, daemon=True)
  57. self._worker.start()
  58. return filename
  59. def write(self, audio: np.ndarray) -> None:
  60. status = self.get_status()
  61. if not status.recording:
  62. return
  63. data = np.asarray(audio, dtype=np.float32)
  64. if data.ndim == 1:
  65. if status.channels != 1:
  66. return
  67. frames = data.shape[0]
  68. int16_data = (np.clip(data, -1.0, 1.0) * 32767.0).astype(np.int16)
  69. payload = int16_data.tobytes()
  70. elif data.ndim == 2:
  71. if data.shape[1] != status.channels:
  72. return
  73. frames = data.shape[0]
  74. int16_data = (np.clip(data, -1.0, 1.0) * 32767.0).astype(np.int16)
  75. payload = int16_data.reshape(-1).tobytes()
  76. else:
  77. return
  78. try:
  79. self._queue.put_nowait(payload)
  80. except queue.Full:
  81. return
  82. with self._lock:
  83. self._frames_written += frames
  84. self._status.duration_sec = self._frames_written / float(status.sample_rate)
  85. def stop(self) -> RecorderStatus:
  86. with self._lock:
  87. if not self._status.recording:
  88. return RecorderStatus(**asdict(self._status))
  89. self._status.recording = False
  90. worker = self._worker
  91. self._stop_event.set()
  92. if worker is not None:
  93. worker.join(timeout=2.0)
  94. with self._lock:
  95. if self._wave_file is not None:
  96. self._wave_file.close()
  97. self._wave_file = None
  98. self._worker = None
  99. self._stop_event.clear()
  100. return RecorderStatus(**asdict(self._status))
  101. def _writer_loop(self) -> None:
  102. while not self._stop_event.is_set() or not self._queue.empty():
  103. try:
  104. payload = self._queue.get(timeout=0.1)
  105. except queue.Empty:
  106. continue
  107. with self._lock:
  108. if self._wave_file is not None:
  109. self._wave_file.writeframes(payload)
  110. def _clear_queue_locked(self) -> None:
  111. while True:
  112. try:
  113. self._queue.get_nowait()
  114. except queue.Empty:
  115. break
  116. def resolve_recording_path(recordings_dir: Path, filename: str) -> Path:
  117. root = Path(recordings_dir).resolve()
  118. candidate = (root / filename).resolve()
  119. if root not in candidate.parents and candidate != root:
  120. raise ValueError("Invalid recording path")
  121. return candidate
  122. def list_recordings(recordings_dir: Path) -> list[dict[str, object]]:
  123. results: list[dict[str, object]] = []
  124. folder = Path(recordings_dir)
  125. folder.mkdir(parents=True, exist_ok=True)
  126. for wav_file in sorted(folder.glob("*.wav"), key=lambda p: p.stat().st_mtime, reverse=True):
  127. try:
  128. with wave.open(str(wav_file), "rb") as wf:
  129. frames = wf.getnframes()
  130. sample_rate = wf.getframerate()
  131. channels = wf.getnchannels()
  132. except (wave.Error, EOFError, OSError):
  133. continue
  134. size_bytes = wav_file.stat().st_size
  135. mtime = datetime.fromtimestamp(wav_file.stat().st_mtime).isoformat(timespec="seconds")
  136. duration_sec = float(frames) / float(sample_rate) if sample_rate else 0.0
  137. results.append(
  138. {
  139. "filename": wav_file.name,
  140. "date": mtime,
  141. "duration_sec": round(duration_sec, 3),
  142. "size_bytes": size_bytes,
  143. "channels": channels,
  144. "sample_rate": sample_rate,
  145. }
  146. )
  147. return results