Subtitle class: get format dinamically (fix #2514)v1.4.4-beta.9

2 files changed, 46 insertions, 10 deletions

diff --git a/custom_libs/subliminal_patch/subtitle.py b/custom_libs/subliminal_patch/subtitle.py
index d0cd585ed..36afea6d2 100644
--- a/custom_libs/subliminal_patch/subtitle.py
+++ b/custom_libs/subliminal_patch/subtitle.py
@@ -62,7 +62,7 @@ class Subtitle(Subtitle_):
     _guessed_encoding = None
     _is_valid = False
     use_original_format = False
-    format = "srt" # default format is srt
+    # format = "srt" # default format is srt
 
     def __init__(self, language, hearing_impaired=False, page_link=None, encoding=None, mods=None, original_format=False):
         # set subtitle language to hi if it's hearing_impaired
@@ -74,6 +74,21 @@ class Subtitle(Subtitle_):
         self.mods = mods
         self._is_valid = False
         self.use_original_format = original_format
+        self._og_format = None
+
+    @property
+    def format(self):
+        if self.use_original_format and self._og_format is not None:
+            logger.debug("Original format requested [%s]", self._og_format)
+            return self._og_format
+
+        logger.debug("Will assume srt format")
+        return "srt"
+
+    # Compatibility
+    @format.setter
+    def format(self, val):
+        self._og_format = val
 
     def __repr__(self):
         r_info = str(self.release_info or "").replace("\n", " | ").strip()
@@ -292,11 +307,13 @@ class Subtitle(Subtitle_):
                     logger.info("Got FPS from MicroDVD subtitle: %s", subs.fps)
                 else:
                     logger.info("Got format: %s", subs.format)
-                    if self.use_original_format:
-                        self.format = subs.format
-                        self._is_valid = True
-                        logger.debug("Using original format")
-                        return True
+                    self._og_format = subs.format
+                    self._is_valid = True
+                    # if self.use_original_format:
+                    #    self.format = subs.format
+                    #    self._is_valid = True
+                    #    logger.debug("Using original format")
+                    return True
 
             except pysubs2.UnknownFPSError:
                 # if parsing failed, use frame rate from provider
diff --git a/tests/subliminal_patch/test_embeddedsubtitles.py b/tests/subliminal_patch/test_embeddedsubtitles.py
index d7f157de1..b7bfaa0a8 100644
--- a/tests/subliminal_patch/test_embeddedsubtitles.py
+++ b/tests/subliminal_patch/test_embeddedsubtitles.py
@@ -8,9 +8,9 @@ from fese.exceptions import LanguageNotFound
 import pytest
 from subliminal_patch.core import Episode
 from subliminal_patch.core import Movie
+from subliminal_patch.providers.embeddedsubtitles import _clean_ass_subtitles
 from subliminal_patch.providers.embeddedsubtitles import (
     _discard_possible_incomplete_subtitles,
-    _clean_ass_subtitles,
 )
 from subliminal_patch.providers.embeddedsubtitles import _get_pretty_release_name
 from subliminal_patch.providers.embeddedsubtitles import _MemoizedFFprobeVideoContainer
@@ -127,7 +127,9 @@ def fake_streams():
 
 @pytest.mark.parametrize("tags_", [{}, {"language": "und", "title": "Unknown"}])
 def test_list_subtitles_unknown_as_fallback(mocker, tags_, video_single_language):
-    with EmbeddedSubtitlesProvider(unknown_as_fallback=True, fallback_lang="en") as provider:
+    with EmbeddedSubtitlesProvider(
+        unknown_as_fallback=True, fallback_lang="en"
+    ) as provider:
         fake = FFprobeSubtitleStream(
             {"index": 3, "codec_name": "subrip", "tags": tags_}
         )
@@ -144,7 +146,9 @@ def test_list_subtitles_unknown_as_fallback(mocker, tags_, video_single_language
 def test_list_subtitles_unknown_as_fallback_w_real_english_subtitles(
     video_single_language, mocker
 ):
-    with EmbeddedSubtitlesProvider(unknown_as_fallback=True, fallback_lang="en") as provider:
+    with EmbeddedSubtitlesProvider(
+        unknown_as_fallback=True, fallback_lang="en"
+    ) as provider:
         fakes = [
             FFprobeSubtitleStream(
                 {"index": 3, "codec_name": "subrip", "tags": {"language": "und"}}
@@ -165,7 +169,7 @@ def test_list_subtitles_unknown_as_fallback_w_real_english_subtitles(
 
 @pytest.mark.parametrize("tags_", [{}, {"language": "und", "title": "Unknown"}])
 def test_list_subtitles_unknown_as_fallback_disabled(tags_):
-    with EmbeddedSubtitlesProvider(unknown_as_fallback=False,fallback_lang="en"):
+    with EmbeddedSubtitlesProvider(unknown_as_fallback=False, fallback_lang="en"):
         with pytest.raises(LanguageNotFound):
             assert FFprobeSubtitleStream(
                 {"index": 3, "codec_name": "subrip", "tags": tags_}
@@ -362,6 +366,21 @@ def test_download_subtitle_single(video_single_language):
         assert subtitle.is_valid()
 
 
+def test_download_subtitle_single_is_ass(video_single_language):
+    with EmbeddedSubtitlesProvider() as provider:
+        subtitle = provider.list_subtitles(
+            video_single_language, {Language.fromalpha2("en")}
+        )[0]
+        provider.download_subtitle(subtitle)
+        assert subtitle.is_valid()
+
+        assert subtitle.format == "srt"
+
+        subtitle.use_original_format = True
+
+        assert subtitle.format == "ass"
+
+
 def test_memoized(video_single_language, mocker):
     with EmbeddedSubtitlesProvider() as provider:
         provider.list_subtitles(video_single_language, {Language.fromalpha2("en")})