From 4f4542c41cf7c44bbed33c6a83dce38aeb611744 Mon Sep 17 00:00:00 2001
From: Andy <scene@scene.com>
Date: Sat, 19 Jul 2025 05:44:33 +0000
Subject: [PATCH] =?UTF-8?q?fixes:=20=E2=9C=A8=20add=20SubtitleCodecChoice?=
 =?UTF-8?q?=20for=20resolving=20issues=20with=20config=20clicktype=20selec?=
 =?UTF-8?q?tion,=20using=20names=20like=20VTT=20or=20SRT=20was=20not=20wor?=
 =?UTF-8?q?king=20as=20expected?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Introduced `SubtitleCodecChoice` to allow selection of subtitle codecs with support for enum names, values, and common aliases.
* Updated `--sub-format` option in `dl.py` to utilize the new `SubtitleCodecChoice`.
---
 unshackle/commands/dl.py            |  4 +-
 unshackle/core/utils/click_types.py | 85 +++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py
index 2ddc515..45a34e3 100644
--- a/unshackle/commands/dl.py
+++ b/unshackle/commands/dl.py
@@ -58,7 +58,7 @@ from unshackle.core.tracks.attachment import Attachment
 from unshackle.core.utilities import get_system_fonts, is_close_match, time_elapsed_since
 from unshackle.core.utils import tags
 from unshackle.core.utils.click_types import (LANGUAGE_RANGE, QUALITY_LIST, SEASON_RANGE, ContextData, MultipleChoice,
-                                              VideoCodecChoice)
+                                              SubtitleCodecChoice, VideoCodecChoice)
 from unshackle.core.utils.collections import merge_dict
 from unshackle.core.utils.subprocess import ffprobe
 from unshackle.core.vaults import Vaults
@@ -179,7 +179,7 @@ class dl:
     )
     @click.option(
         "--sub-format",
-        type=click.Choice(Subtitle.Codec, case_sensitive=False),
+        type=SubtitleCodecChoice(Subtitle.Codec),
         default="srt",
         help="Set Output Subtitle Format, only converting if necessary.",
     )
diff --git a/unshackle/core/utils/click_types.py b/unshackle/core/utils/click_types.py
index 479dcbf..8707732 100644
--- a/unshackle/core/utils/click_types.py
+++ b/unshackle/core/utils/click_types.py
@@ -42,6 +42,91 @@ class VideoCodecChoice(click.Choice):
         self.fail(f"'{value}' is not a valid video codec", param, ctx)
 
 
+class SubtitleCodecChoice(click.Choice):
+    """
+    A custom Choice type for subtitle codecs that accepts both enum names, values, and common aliases.
+
+    Accepts:
+    - Enum names: subrip, substationalpha, substationalphav4, timedtextmarkuplang, webvtt, ftml, fvtt
+    - Enum values: SRT, SSA, ASS, TTML, VTT, STPP, WVTT
+    - Common aliases: srt (for SubRip)
+    """
+
+    def __init__(self, codec_enum):
+        self.codec_enum = codec_enum
+        # Build choices from enum names, values, and common aliases
+        choices = []
+        aliases = {}
+
+        for codec in codec_enum:
+            choices.append(codec.name.lower())  # e.g., "subrip", "webvtt"
+
+            # Only add the value if it's different from common aliases
+            value_lower = codec.value.lower()
+
+            # Add common aliases and track them
+            if codec.name == "SubRip":
+                if "srt" not in choices:
+                    choices.append("srt")
+                aliases["srt"] = codec
+            elif codec.name == "WebVTT":
+                if "vtt" not in choices:
+                    choices.append("vtt")
+                aliases["vtt"] = codec
+                # Also add the enum value if different
+                if value_lower != "vtt" and value_lower not in choices:
+                    choices.append(value_lower)
+            elif codec.name == "SubStationAlpha":
+                if "ssa" not in choices:
+                    choices.append("ssa")
+                aliases["ssa"] = codec
+                # Also add the enum value if different
+                if value_lower != "ssa" and value_lower not in choices:
+                    choices.append(value_lower)
+            elif codec.name == "SubStationAlphav4":
+                if "ass" not in choices:
+                    choices.append("ass")
+                aliases["ass"] = codec
+                # Also add the enum value if different
+                if value_lower != "ass" and value_lower not in choices:
+                    choices.append(value_lower)
+            elif codec.name == "TimedTextMarkupLang":
+                if "ttml" not in choices:
+                    choices.append("ttml")
+                aliases["ttml"] = codec
+                # Also add the enum value if different
+                if value_lower != "ttml" and value_lower not in choices:
+                    choices.append(value_lower)
+            else:
+                # For other codecs, just add the enum value
+                if value_lower not in choices:
+                    choices.append(value_lower)
+
+        self.aliases = aliases
+        super().__init__(choices, case_sensitive=False)
+
+    def convert(self, value: Any, param: Optional[click.Parameter] = None, ctx: Optional[click.Context] = None):
+        if not value:
+            return None
+
+        # First try to convert using the parent class
+        converted_value = super().convert(value, param, ctx)
+
+        # Check aliases first
+        if converted_value.lower() in self.aliases:
+            return self.aliases[converted_value.lower()]
+
+        # Now map the converted value back to the enum
+        for codec in self.codec_enum:
+            if converted_value.lower() == codec.name.lower():
+                return codec
+            if converted_value.lower() == codec.value.lower():
+                return codec
+
+        # This shouldn't happen if the parent conversion worked
+        self.fail(f"'{value}' is not a valid subtitle codec", param, ctx)
+
+
 class ContextData:
     def __init__(self, config: dict, cdm: WidevineCdm, proxy_providers: list, profile: Optional[str] = None):
         self.config = config