From 33ab54a626d7ff8520b699d79d3ae5383d6e255f Mon Sep 17 00:00:00 2001
From: Mathieu Virbel <mat@meltingrocks.com>
Date: Wed, 16 Aug 2023 10:40:40 +0200
Subject: [PATCH] server: replace wave module with pyav directly

Closes #87
---
 .../reflector/processors/audio_file_writer.py | 26 +++++++++++--------
 server/reflector/processors/audio_merge.py    | 18 ++++++++-----
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/server/reflector/processors/audio_file_writer.py b/server/reflector/processors/audio_file_writer.py
index c597f81d..d67db65e 100644
--- a/server/reflector/processors/audio_file_writer.py
+++ b/server/reflector/processors/audio_file_writer.py
@@ -1,6 +1,5 @@
 from reflector.processors.base import Processor
 import av
-import wave
 from pathlib import Path
 
 
@@ -17,19 +16,24 @@ class AudioFileWriterProcessor(Processor):
         if isinstance(path, str):
             path = Path(path)
         self.path = path
-        self.fd = None
+        self.out_container = None
+        self.out_stream = None
 
     async def _push(self, data: av.AudioFrame):
-        if not self.fd:
+        if not self.out_container:
             self.path.parent.mkdir(parents=True, exist_ok=True)
-            self.fd = wave.open(self.path.as_posix(), "wb")
-            self.fd.setnchannels(len(data.layout.channels))
-            self.fd.setsampwidth(data.format.bytes)
-            self.fd.setframerate(data.sample_rate)
-        self.fd.writeframes(data.to_ndarray().tobytes())
+            self.out_container = av.open(self.path.as_posix(), "w", format="wav")
+            self.out_stream = self.out_container.add_stream(
+                "pcm_s16le", rate=data.sample_rate
+            )
+            for packet in self.out_stream.encode(data):
+                self.out_container.mux(packet)
         await self.emit(data)
 
     async def _flush(self):
-        if self.fd:
-            self.fd.close()
-            self.fd = None
+        if self.out_container:
+            for packet in self.out_stream.encode(None):
+                self.out_container.mux(packet)
+            self.out_container.close()
+            self.out_container = None
+            self.out_stream = None
diff --git a/server/reflector/processors/audio_merge.py b/server/reflector/processors/audio_merge.py
index 37734a53..34c1741e 100644
--- a/server/reflector/processors/audio_merge.py
+++ b/server/reflector/processors/audio_merge.py
@@ -3,7 +3,6 @@ from reflector.processors.types import AudioFile
 from time import monotonic_ns
 from uuid import uuid4
 import io
-import wave
 import av
 
 
@@ -28,12 +27,16 @@ class AudioMergeProcessor(Processor):
         # create audio file
         uu = uuid4().hex
         fd = io.BytesIO()
-        with wave.open(fd, "wb") as wf:
-            wf.setnchannels(channels)
-            wf.setsampwidth(sample_width)
-            wf.setframerate(sample_rate)
-            for frame in data:
-                wf.writeframes(frame.to_ndarray().tobytes())
+
+        out_container = av.open(fd, "w", format="wav")
+        out_stream = out_container.add_stream("pcm_s16le", rate=sample_rate)
+        for frame in data:
+            for packet in out_stream.encode(frame):
+                out_container.mux(packet)
+        for packet in out_stream.encode(None):
+            out_container.mux(packet)
+        out_container.close()
+        fd.seek(0)
 
         # emit audio file
         audiofile = AudioFile(
@@ -44,4 +47,5 @@ class AudioMergeProcessor(Processor):
             sample_width=sample_width,
             timestamp=data[0].pts * data[0].time_base,
         )
+
         await self.emit(audiofile)