Handle InvalidFile with Embedded Subtitles Provider (close #1652)

2 files changed, 25 insertions, 7 deletions

diff --git a/libs/subliminal_patch/providers/embeddedsubtitles.py b/libs/subliminal_patch/providers/embeddedsubtitles.py
index 4f3a5dfb5..4ac49a7b8 100644
--- a/libs/subliminal_patch/providers/embeddedsubtitles.py
+++ b/libs/subliminal_patch/providers/embeddedsubtitles.py
@@ -8,12 +8,14 @@ import tempfile
 from babelfish import language_converters
 import fese
 from fese import check_integrity
+from fese import InvalidFile
 from fese import FFprobeSubtitleStream
 from fese import FFprobeVideoContainer
 from fese import to_srt
 from subliminal.subtitle import fix_line_ending
 from subliminal_patch.core import Episode
 from subliminal_patch.core import Movie
+from subliminal_patch.exceptions import MustGetBlacklisted
 from subliminal_patch.providers import Provider
 from subliminal_patch.subtitle import Subtitle
 from subzero.language import Language
@@ -28,7 +30,7 @@ class EmbeddedSubtitle(Subtitle):
     provider_name = "embeddedsubtitles"
     hash_verifiable = False
 
-    def __init__(self, stream, container, matches):
+    def __init__(self, stream, container, matches, media_type):
         super().__init__(stream.language, stream.disposition.hearing_impaired)
         if stream.disposition.forced:
             self.language = Language.rebuild(stream.language, forced=True)
@@ -36,9 +38,11 @@ class EmbeddedSubtitle(Subtitle):
         self.stream: FFprobeSubtitleStream = stream
         self.container: FFprobeVideoContainer = container
         self.forced = stream.disposition.forced
-        self._matches: set = matches
         self.page_link = self.container.path
         self.release_info = os.path.basename(self.page_link)
+        self.media_type = media_type
+
+        self._matches: set = matches
 
     def get_matches(self, video):
         if self.hearing_impaired:
@@ -97,7 +101,7 @@ class EmbeddedSubtitlesProvider(Provider):
         # Remove leftovers
         shutil.rmtree(self._cache_dir, ignore_errors=True)
 
-    def query(self, path: str, languages):
+    def query(self, path: str, languages, media_type):
         video = FFprobeVideoContainer(path)
 
         try:
@@ -144,14 +148,21 @@ class EmbeddedSubtitlesProvider(Provider):
         if self._hi_fallback:
             _check_hi_fallback(allowed_streams, languages)
 
-        return [EmbeddedSubtitle(stream, video, {"hash"}) for stream in allowed_streams]
+        return [
+            EmbeddedSubtitle(stream, video, {"hash"}, media_type)
+            for stream in allowed_streams
+        ]
 
     def list_subtitles(self, video, languages):
         if not os.path.isfile(video.original_path):
             logger.debug("Ignoring inexistent file: %s", video.original_path)
             return []
 
-        return self.query(video.original_path, languages)
+        return self.query(
+            video.original_path,
+            languages,
+            "series" if isinstance(video, Episode) else "movie",
+        )
 
     def download_subtitle(self, subtitle):
         path = self._get_subtitle_path(subtitle)
@@ -175,7 +186,10 @@ class EmbeddedSubtitlesProvider(Provider):
         # Get the subtitle file by index
         subtitle_path = cached_path[subtitle.stream.index]
 
-        check_integrity(subtitle.stream, subtitle_path)
+        try:
+            check_integrity(subtitle.stream, subtitle_path)
+        except InvalidFile as error:
+            raise MustGetBlacklisted(subtitle.id, subtitle.media_type) from error
 
         # Convert to SRT if the subtitle is ASS
         new_subtitle_path = to_srt(subtitle_path, remove_source=True)
diff --git a/tests/subliminal_patch/test_embeddedsubtitles.py b/tests/subliminal_patch/test_embeddedsubtitles.py
index c9c8e7cd5..e42330425 100644
--- a/tests/subliminal_patch/test_embeddedsubtitles.py
+++ b/tests/subliminal_patch/test_embeddedsubtitles.py
@@ -6,6 +6,7 @@ import pytest
 import fese
 from fese import FFprobeSubtitleStream
 from subliminal_patch.core import Episode, Movie
+from subliminal_patch.exceptions import MustGetBlacklisted
 from subliminal_patch.providers.embeddedsubtitles import EmbeddedSubtitlesProvider
 from subzero.language import Language
 
@@ -223,5 +224,8 @@ def test_download_invalid_subtitle(video_single_language):
         provider._cached_paths[subtitle.container.path] = {
             subtitle.stream.index: "dummy.srt"
         }
-        with pytest.raises(fese.InvalidFile):
+        try:
             provider.download_subtitle(subtitle)
+        except MustGetBlacklisted as error:
+            assert error.id == subtitle.id
+            assert error.media_type == subtitle.media_type