From f1d2b20ab71a2aa2fef48168c534cea6f40512f1 Mon Sep 17 00:00:00 2001
From: WyattBlue <wyattblue@auto-editor.com>
Date: Mon, 16 Sep 2024 04:46:17 -0400
Subject: [PATCH] Use PyAV to make wav files

---
 auto_editor/edit.py      |  2 +-
 auto_editor/ffwrapper.py | 25 +++++++++++++++++++++++++
 auto_editor/timeline.py  |  8 ++++----
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/auto_editor/edit.py b/auto_editor/edit.py
index cea5fe016..da4987a07 100644
--- a/auto_editor/edit.py
+++ b/auto_editor/edit.py
@@ -240,7 +240,7 @@ def edit_media(paths: list[str], ffmpeg: FFmpeg, args: Args, log: Log) -> None:
         if is_resolve:
             from auto_editor.timeline import set_stream_to_0
 
-            set_stream_to_0(tl, ffmpeg, log)
+            set_stream_to_0(tl, log)
 
         fcp11_write_xml(export_ops["name"], output, is_resolve, tl, log)
         return
diff --git a/auto_editor/ffwrapper.py b/auto_editor/ffwrapper.py
index 2b512fb71..7ec8f856f 100644
--- a/auto_editor/ffwrapper.py
+++ b/auto_editor/ffwrapper.py
@@ -127,6 +127,31 @@ def pipe(self, cmd: list[str]) -> str:
         return output
 
 
+def mux(input: Path, output: Path, stream: int, codec: str | None = None) -> None:
+    input_container = av.open(input, "r")
+    output_container = av.open(output, "w")
+
+    input_audio_stream = input_container.streams.audio[stream]
+
+    if codec is None:
+        codec = "pcm_s16le"
+
+    output_audio_stream = output_container.add_stream(codec)
+    assert isinstance(output_audio_stream, av.audio.AudioStream)
+
+    for frame in input_container.decode(input_audio_stream):
+        packet = output_audio_stream.encode(frame)
+        if packet:
+            output_container.mux(packet)
+
+    packet = output_audio_stream.encode(None)
+    if packet:
+        output_container.mux(packet)
+
+    output_container.close()
+    input_container.close()
+
+
 @dataclass(slots=True, frozen=True)
 class VideoStream:
     width: int
diff --git a/auto_editor/timeline.py b/auto_editor/timeline.py
index 195ef50aa..b5d680a82 100644
--- a/auto_editor/timeline.py
+++ b/auto_editor/timeline.py
@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from auto_editor.ffwrapper import initFileInfo
+from auto_editor.ffwrapper import initFileInfo, mux
 from auto_editor.lib.contracts import *
 from auto_editor.utils.cmdkw import Required, pAttr, pAttrs
 from auto_editor.utils.types import color, natural, number, threshold
@@ -14,7 +14,7 @@
     from pathlib import Path
     from typing import Any
 
-    from auto_editor.ffwrapper import FFmpeg, FileInfo
+    from auto_editor.ffwrapper import FileInfo
     from auto_editor.utils.chunks import Chunks
     from auto_editor.utils.log import Log
 
@@ -303,7 +303,7 @@ def make_tracks_dir(path: Path) -> Path:
     return tracks_dir
 
 
-def set_stream_to_0(tl: v3, ffmpeg: FFmpeg, log: Log) -> None:
+def set_stream_to_0(tl: v3, log: Log) -> None:
     src = tl.src
     assert src is not None
     fold = make_tracks_dir(src.path)
@@ -312,7 +312,7 @@ def set_stream_to_0(tl: v3, ffmpeg: FFmpeg, log: Log) -> None:
     def make_track(i: int, path: Path) -> FileInfo:
         newtrack = fold / f"{path.stem}_{i}.wav"
         if newtrack not in cache:
-            ffmpeg.run(["-i", f"{path}", "-map", f"0:a:{i}", f"{newtrack}"])
+            mux(path, output=newtrack, stream=i)
             cache[newtrack] = initFileInfo(f"{newtrack}", log)
         return cache[newtrack]