From 0588f47837d9cd8de86fee6a90b4c38ee4287aff Mon Sep 17 00:00:00 2001
From: Maveno <javanaut@refulgent.de>
Date: Sun, 17 Nov 2024 17:57:36 +0100
Subject: [PATCH] =?UTF-8?q?#398=20Default=20/=20Forced=20f=C3=BCr=20unmuxe?=
 =?UTF-8?q?d/imported=20Streams?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/ffx/ffx.py               |  8 ++++++--
 src/ffx/media_descriptor.py  | 17 +++++++++++++++--
 src/ffx/track_codec.py       |  2 +-
 src/ffx/track_descriptor.py  |  3 +++
 src/ffx/track_disposition.py |  8 ++++++++
 5 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/ffx/ffx.py b/src/ffx/ffx.py
index 4bec8cd..4a9db70 100755
--- a/src/ffx/ffx.py
+++ b/src/ffx/ffx.py
@@ -170,7 +170,7 @@ def unmux(ctx,
 
             season = fp.getSeason()
             episode = fp.getEpisode()
-    
+
             #TODO: Recognition für alle Formate anpassen
             targetLabel = label if label else fp.getFileBasename()
             targetIndicator = f"_S{season}E{episode}" if label and season != -1 and episode != -1 else ''
@@ -185,9 +185,13 @@ def unmux(ctx,
                 
                 if trackDescriptor.getType() == TrackType.SUBTITLE or not subtitles_only:
 
-                    # SEASON_EPISODE_STREAM_LANGUAGE_MATCH = '[sS]([0-9]+)[eE]([0-9]+)_([0-9]+)_([a-z]{3})'
+                    # SEASON_EPISODE_STREAM_LANGUAGE_MATCH = '[sS]([0-9]+)[eE]([0-9]+)_([0-9]+)_([a-z]{3})(?:_([A-Z]{3}))*'
                     targetPrefix = f"{targetLabel}{targetIndicator}_{trackDescriptor.getIndex()}_{trackDescriptor.getLanguage().threeLetter()}"
 
+                    td: TrackDisposition
+                    for td in sorted(trackDescriptor.getDispositionSet(), key=lambda d: d.index()):
+                        targetPrefix += f"_{td.indicator()}"
+
                     unmuxSequence = getUnmuxSequence(trackDescriptor, sourcePath, targetPrefix, targetDirectory = output_directory)
                     
                     if unmuxSequence:
diff --git a/src/ffx/media_descriptor.py b/src/ffx/media_descriptor.py
index 068173a..619046c 100644
--- a/src/ffx/media_descriptor.py
+++ b/src/ffx/media_descriptor.py
@@ -31,7 +31,7 @@ class MediaDescriptor:
     #407 remove as well
     EXCLUDED_MEDIA_TAGS = ["creation_time"]
 
-    SEASON_EPISODE_STREAM_LANGUAGE_MATCH = '[sS]([0-9]+)[eE]([0-9]+)_([0-9]+)_([a-z]{3})'
+    SEASON_EPISODE_STREAM_LANGUAGE_MATCH = '[sS]([0-9]+)[eE]([0-9]+)_([0-9]+)_([a-z]{3})(?:_([A-Z]{3}))*'
     SUBTITLE_FILE_EXTENSION = 'vtt'
 
     def __init__(self, **kwargs):
@@ -404,7 +404,7 @@ class MediaDescriptor:
             if importedFilePath:
 
                 self.__logger.info(f"Substituting subtitle stream #{td.getIndex()} "
-                                    + f"({td.getType().label()}:{td.getSubIndex()})"
+                                    + f"({td.getType().label()}:{td.getSubIndex()}) "
                                     + f"with import from file {td.getExternalSourceFilePath()}")
 
                 importFileTokens += [
@@ -483,6 +483,16 @@ class MediaDescriptor:
                         subtitleFileDescriptor["index"] = int(sesl_result.group(3))
                         subtitleFileDescriptor["language"] = sesl_result.group(4)
 
+                        dispSet = set()
+                        dispCaptGroups = sesl_result.groups()
+                        numCaptGroups = len(dispCaptGroups)
+                        if numCaptGroups > 4:
+                            for groupIndex in range(numCaptGroups - 4):
+                                disp = TrackDisposition.fromIndicator(dispCaptGroups[groupIndex + 4])
+                                if disp is not None:
+                                    dispSet.add(disp)
+                        subtitleFileDescriptor["disposition_set"] = dispSet
+
                         subtitleFileDescriptors.append(subtitleFileDescriptor)
 
         self.__logger.debug(f"searchSubtitleFiles(): Available subtitle files {subtitleFileDescriptors}")
@@ -525,6 +535,9 @@ class MediaDescriptor:
                 self.__logger.debug(f"importSubtitles(): Found matching subtitle file {msfd['path']}")
                 matchingSubtitleTrackDescriptor[0].setExternalSourceFilePath(msfd["path"])
 
+                # TODO: Check if useful
+                # matchingSubtitleTrackDescriptor[0].setDispositionSet(msfd["disposition_set"])
+
 
     def getConfiguration(self, label: str = ''):
         yield f"--- {label if label else 'MediaDescriptor '+str(id(self))} {' '.join([str(k)+'='+str(v) for k,v in self.__mediaTags.items()])}"
diff --git a/src/ffx/track_codec.py b/src/ffx/track_codec.py
index 8539ab4..2de8ab7 100644
--- a/src/ffx/track_codec.py
+++ b/src/ffx/track_codec.py
@@ -28,7 +28,7 @@ class TrackCodec(Enum):
 
     def extension(self):
         """Returns the corresponding extension"""
-        return int(self.value['extension'])
+        return str(self.value['extension'])
 
     @staticmethod
     def identify(identifier: str):
diff --git a/src/ffx/track_descriptor.py b/src/ffx/track_descriptor.py
index 5c24006..f45a332 100644
--- a/src/ffx/track_descriptor.py
+++ b/src/ffx/track_descriptor.py
@@ -309,6 +309,9 @@ class TrackDescriptor:
     def getDispositionSet(self):
         return self.__dispositionSet
 
+    def setDispositionSet(self, dispositionSet: set):
+        self.__dispositionSet = dispositionSet
+
     def getDispositionFlag(self, disposition: TrackDisposition) -> bool:
         return bool(disposition in self.__dispositionSet)
 
diff --git a/src/ffx/track_disposition.py b/src/ffx/track_disposition.py
index 4cec46d..3d2709e 100644
--- a/src/ffx/track_disposition.py
+++ b/src/ffx/track_disposition.py
@@ -66,3 +66,11 @@ class TrackDisposition(Enum):
             return matchingDispositions[0]
         else:
             return None
+
+    @staticmethod
+    def fromIndicator(indicator: str):
+        matchingDispositions = [d for d in TrackDisposition if d.indicator() == str(indicator)]
+        if matchingDispositions:
+            return matchingDispositions[0]
+        else:
+            return None