From 2cee35066110b0e9b715670f39792dde124f857d Mon Sep 17 00:00:00 2001
From: Chris Busillo <info@shinycomputers.com>
Date: Sun, 23 Jun 2024 12:06:15 -0400
Subject: [PATCH] refactor subs to own main function and add stage

---
 bd_to_avp/gui/main_window.py   |  2 +-
 bd_to_avp/modules/config.py    |  6 ++++--
 bd_to_avp/modules/container.py | 12 ++++--------
 bd_to_avp/modules/process.py   |  8 ++++----
 bd_to_avp/modules/sub.py       |  9 ++++++++-
 5 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/bd_to_avp/gui/main_window.py b/bd_to_avp/gui/main_window.py
index 83b00fd..128bee6 100644
--- a/bd_to_avp/gui/main_window.py
+++ b/bd_to_avp/gui/main_window.py
@@ -253,7 +253,7 @@ def handle_mkv_creation_error(self, error: MKVCreationError) -> None:
 
         if result == QMessageBox.StandardButton.Yes:
             config.continue_on_error = True
-            config.start_stage = Stage.EXTRACT_MVC_AUDIO_AND_SUB
+            config.start_stage = Stage.EXTRACT_MVC_AND_AUDIO
             self.start_processing(is_continuing=True)
             return
 
diff --git a/bd_to_avp/modules/config.py b/bd_to_avp/modules/config.py
index d799ac8..c959a94 100644
--- a/bd_to_avp/modules/config.py
+++ b/bd_to_avp/modules/config.py
@@ -12,7 +12,8 @@
 
 class Stage(Enum):
     CREATE_MKV = auto()
-    EXTRACT_MVC_AUDIO_AND_SUB = auto()
+    EXTRACT_MVC_AND_AUDIO = auto()
+    EXTRACT_SUBTITLES = auto()
     CREATE_LEFT_RIGHT_FILES = auto()
     COMBINE_TO_MV_HEVC = auto()
     UPSCALE_VIDEO = auto()
@@ -26,7 +27,8 @@ def __str__(self) -> str:
     def human_readable(self) -> str:
         return {
             "CREATE_MKV": "Create MKV",
-            "EXTRACT_MVC_AUDIO_AND_SUB": "Extract MVC Audio and Sub",
+            "EXTRACT_MVC_AND_AUDIO": "Extract MVC and Audio",
+            "EXTRACT_SUBTITLES": "Extract Subtitles",
             "CREATE_LEFT_RIGHT_FILES": "Create Left Right Files",
             "UPSCALE_VIDEO": "Upscale Video",
             "COMBINE_TO_MV_HEVC": "Combine to MV HEVC",
diff --git a/bd_to_avp/modules/container.py b/bd_to_avp/modules/container.py
index 3faef5d..115d6c1 100644
--- a/bd_to_avp/modules/container.py
+++ b/bd_to_avp/modules/container.py
@@ -5,7 +5,6 @@
 from babelfish import Language
 
 from bd_to_avp.modules.config import Stage, config
-from bd_to_avp.modules.sub import extract_subtitle_to_srt
 from bd_to_avp.modules.command import run_command, run_ffmpeg_print_errors
 from bd_to_avp.modules.util import sorted_files_by_creation_filtered_on_suffix
 
@@ -42,18 +41,15 @@ def create_muxed_file(
     return muxed_path
 
 
-def create_mvc_audio_and_subtitle_files(
+def create_mvc_and_audio(
     disc_name: str,
-    mkv_output_path: Path | None,
+    mkv_output_path: Path,
     output_folder: Path,
 ) -> tuple[Path, Path]:
     video_output_path = output_folder / f"{disc_name}_mvc.h264"
     audio_output_path = output_folder / f"{disc_name}_audio_PCM.mov"
 
-    if config.start_stage.value <= Stage.EXTRACT_MVC_AUDIO_AND_SUB.value and mkv_output_path:
-        if not config.skip_subtitles:
-            extract_subtitle_to_srt(mkv_output_path, output_folder)
-
+    if config.start_stage.value <= Stage.EXTRACT_MVC_AND_AUDIO.value:
         extract_mvc_and_audio(
             mkv_output_path,
             video_output_path,
@@ -113,7 +109,7 @@ def mux_video_audio_subs(mv_hevc_path: Path, audio_path: Path, muxed_path: Path,
         output_track_index += 1
 
     command += [muxed_path]
-    run_command(command, "Mux video, audio, and subtitles.")
+    run_command(command, "mux video, audio, and subtitles.")
 
 
 def get_audio_stream_data(file_path: Path) -> list[dict[str, Any]]:
diff --git a/bd_to_avp/modules/process.py b/bd_to_avp/modules/process.py
index ecf5a66..c4fe719 100755
--- a/bd_to_avp/modules/process.py
+++ b/bd_to_avp/modules/process.py
@@ -4,7 +4,7 @@
 
 from bd_to_avp.modules.audio import create_transcoded_audio_file
 from bd_to_avp.modules.config import config
-from bd_to_avp.modules.container import create_muxed_file, create_mvc_audio_and_subtitle_files
+from bd_to_avp.modules.container import create_muxed_file, create_mvc_and_audio
 from bd_to_avp.modules.disc import create_mkv_file, get_disc_and_mvc_video_info
 from bd_to_avp.modules.file import (
     file_exists_normalized,
@@ -12,6 +12,7 @@
     prepare_output_folder_for_source,
     remove_folder_if_exists,
 )
+from bd_to_avp.modules.sub import create_srt_from_mkv
 from bd_to_avp.modules.video import (
     create_left_right_files,
     create_mv_hevc_file,
@@ -64,9 +65,8 @@ def process_each() -> None:
 
     mkv_output_path = create_mkv_file(output_folder, disc_info, config.language_code)
     crop_params = detect_crop_parameters(mkv_output_path)
-    audio_output_path, video_output_path = create_mvc_audio_and_subtitle_files(
-        disc_info.name, mkv_output_path, output_folder
-    )
+    audio_output_path, video_output_path = create_mvc_and_audio(disc_info.name, mkv_output_path, output_folder)
+    create_srt_from_mkv(mkv_output_path, output_folder)
     left_output_path, right_output_path = create_left_right_files(
         disc_info, output_folder, video_output_path, crop_params
     )
diff --git a/bd_to_avp/modules/sub.py b/bd_to_avp/modules/sub.py
index 3b4d119..b696c48 100644
--- a/bd_to_avp/modules/sub.py
+++ b/bd_to_avp/modules/sub.py
@@ -7,7 +7,7 @@
 from babelfish import Language
 from pgsrip import Mkv, Options, pgsrip
 
-from bd_to_avp.modules.config import config
+from bd_to_avp.modules.config import config, Stage
 from bd_to_avp.modules.command import Spinner
 
 
@@ -15,6 +15,13 @@ class SRTCreationError(Exception):
     pass
 
 
+def create_srt_from_mkv(mkv_path: Path, output_path: Path) -> None:
+    if config.start_stage.value <= Stage.EXTRACT_SUBTITLES.value:
+        if config.skip_subtitles:
+            return None
+        extract_subtitle_to_srt(mkv_path, output_path)
+
+
 def extract_subtitle_to_srt(mkv_path: Path, output_path: Path) -> None:
     if config.skip_subtitles:
         return None