Fix MKV crash and read the source audio bitrate from the first track

ffmpeg4discord/twopass.py

@@ -66,6 +66,9 @@
 # Characters that cmd and PowerShell both leave alone, so an argument made only of these needs no quotes.
 WINDOWS_SAFE_CHARS = set(string.ascii_letters + string.digits + "-_./\\:=+")
 
+# The audio bitrate in bits per second to use when the input doesn't report one. It matches the -a default.
+FALLBACK_AUDIO_BR = 96_000
+
 # ffprobe `color_transfer` values that mean the video is HDR: PQ (HDR10, Dolby Vision) and HLG.
 HDR_TRANSFERS = {"smpte2084", "arib-std-b67"}
 
@@ -316,13 +319,8 @@ def __init__(self, filename: Path) -> None:
         self._process_probe()
 
     def _parse_streams(self):
-        """Helper to parse video and audio streams from the probe output.
-
-        Returns:
-            The ffprobe "index" of the *first audio stream* (if any). This is only used for audio bitrate lookup.
-        """
+        """Helper to parse video and audio streams from the probe output."""
 
-        audio_stream = None
         for stream in self.probe["streams"]:
             codec_type = stream["codec_type"]
             if codec_type == "video":
@@ -338,11 +336,8 @@ def _parse_streams(self):
                 framerate_ratio: str = stream.get("r_frame_rate")
                 self.init_framerate = round(int(framerate_ratio.split("/")[0]) / int(framerate_ratio.split("/")[1]))
             elif codec_type == "audio":
-                audio_stream = stream["index"]
                 self.audio_streams.append(stream)
 
-        return audio_stream
-
     def _process_probe(self):
         """
         Processes the ffmpeg probe output to set the duration, size, ratio, framerate, and audio bitrate attributes.
@@ -351,11 +346,12 @@ def _process_probe(self):
         # in mebibytes (MiB, 1 MiB = 1,048,576 bytes), like the target size
         self.input_filesize = float(self.probe["format"]["size"]) * FILE_SIZE_MULT
 
-        audio_stream = self._parse_streams()
+        self._parse_streams()
 
         self.source_audio_br = None
-        if audio_stream is not None:
-            self.source_audio_br = float(self.probe["streams"][audio_stream]["bit_rate"])
+        if self.audio_streams:
+            # MKV files, like OBS recordings, don't report a bitrate for each stream
+            self.source_audio_br = float(self.audio_streams[0].get("bit_rate", FALLBACK_AUDIO_BR))
         else:
             logging.warning("No audio stream found in the media file.")
 

tests/test_twopass.py

@@ -141,6 +141,18 @@ def test_audio_br_defaults_to_the_source_bitrate(make_twopass):
     assert make_twopass()._audio_bitrate(make_settings()) == 128000
 
 
+def test_source_bitrate_comes_from_the_first_audio_track(probe, make_twopass):
+    # e.g. a ShadowPlay clip with game audio first and a lower-bitrate mic track second
+    probe["streams"].append({"codec_type": "audio", "bit_rate": "64000", "index": 2})
+    assert make_twopass().source_audio_br == 128000
+
+
+def test_source_bitrate_falls_back_when_the_input_has_none(probe, make_twopass):
+    # MKV files don't report a bitrate for each stream
+    del probe["streams"][1]["bit_rate"]
+    assert make_twopass().source_audio_br == 96000
+
+
 def test_warns_when_no_audio_stream(probe, make_twopass, caplog):
     probe["streams"].pop()
     make_twopass()