From 9b74ad8610f35275e76b9374643d768c6ef8df53 Mon Sep 17 00:00:00 2001 From: Vitiko Date: Tue, 14 Dec 2021 15:50:45 -0400 Subject: [PATCH] Update Embedded Subtitles provider: add forced and ass-only support --- bazarr/config.py | 1 + bazarr/get_providers.py | 1 + frontend/src/Settings/Providers/list.ts | 6 ++- .../providers/embeddedsubtitles.py | 37 ++++++++++++----- .../test_embeddedsubtitles.py | 40 +++++++++++++++++++ 5 files changed, 73 insertions(+), 12 deletions(-) diff --git a/bazarr/config.py b/bazarr/config.py index c27931344..1982ebaf9 100644 --- a/bazarr/config.py +++ b/bazarr/config.py @@ -192,6 +192,7 @@ defaults = { }, 'embeddedsubtitles': { 'include_ass': 'True', + 'include_srt': 'True', }, 'subsync': { 'use_subsync': 'False', diff --git a/bazarr/get_providers.py b/bazarr/get_providers.py index 9e1d652f6..fd6f75dea 100644 --- a/bazarr/get_providers.py +++ b/bazarr/get_providers.py @@ -201,6 +201,7 @@ def get_providers_auth(): }, 'embeddedsubtitles': { 'include_ass': settings.embeddedsubtitles.getboolean('include_ass'), + 'include_srt': settings.embeddedsubtitles.getboolean('include_srt'), 'cache_dir': os.path.join(args.config_dir, "cache"), 'ffprobe_path': get_binary("ffprobe"), 'ffmpeg_path': get_binary("ffmpeg"), diff --git a/frontend/src/Settings/Providers/list.ts b/frontend/src/Settings/Providers/list.ts index c5defab07..60b0df56d 100644 --- a/frontend/src/Settings/Providers/list.ts +++ b/frontend/src/Settings/Providers/list.ts @@ -51,10 +51,14 @@ export const ProviderList: Readonly = [ name: "Embedded Subtitles", description: "Embedded Subtitles from your Media Files", defaultKey: { + include_srt: true, include_ass: true, }, + message: + "Warning for cloud users: this provider needs to read the entire file in order to extract subtitles.", keyNameOverride: { - include_ass: "Convert embedded ASS to SRT", + include_srt: "Include SRT", + include_ass: "Include ASS (will be converted to SRT)", }, }, { diff --git a/libs/subliminal_patch/providers/embeddedsubtitles.py b/libs/subliminal_patch/providers/embeddedsubtitles.py index bf89340c5..3d9e64959 100644 --- a/libs/subliminal_patch/providers/embeddedsubtitles.py +++ b/libs/subliminal_patch/providers/embeddedsubtitles.py @@ -32,6 +32,7 @@ class EmbeddedSubtitle(Subtitle): super().__init__(stream.language, stream.disposition.hearing_impaired) 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) @@ -55,17 +56,21 @@ class EmbeddedSubtitlesProvider(Provider): Language.fromalpha2(l) for l in language_converters["alpha2"].codes } languages.update(set(Language.rebuild(lang, hi=True) for lang in languages)) - - # TODO: add forced support - # languages.update(set(Language.rebuild(lang, forced=True) for lang in languages)) + languages.update(set(Language.rebuild(lang, forced=True) for lang in languages)) video_types = (Episode, Movie) subtitle_class = EmbeddedSubtitle def __init__( - self, include_ass=True, cache_dir=None, ffprobe_path=None, ffmpeg_path=None + self, + include_ass=True, + include_srt=True, + cache_dir=None, + ffprobe_path=None, + ffmpeg_path=None, ): self._include_ass = include_ass + self._include_srt = include_srt self._cache_dir = os.path.join( cache_dir or tempfile.gettempdir(), self.__class__.__name__.lower() ) @@ -91,7 +96,7 @@ class EmbeddedSubtitlesProvider(Provider): video = FFprobeVideoContainer(path) try: - streams = video.get_subtitles() + streams = filter(_check_allowed_extensions, video.get_subtitles()) except fese.InvalidSource as error: logger.error("Error trying to get subtitles for %s: %s", video, error) streams = [] @@ -99,23 +104,33 @@ class EmbeddedSubtitlesProvider(Provider): if not streams: logger.debug("No subtitles found for container: %s", video) + only_forced = all(lang.forced for lang in languages) + also_forced = any(lang.forced for lang in languages) + subtitles = [] for stream in streams: - # Only subrip and ass are currently supported - if stream.codec_name not in ("subrip", "ass"): - logger.debug("Ignoring codec: %s", stream) + if not self._include_ass and stream.extension == "ass": + logger.debug("Ignoring ASS: %s", stream) continue - if not self._include_ass and stream.codec_name == "ass": - logger.debug("Ignoring ASS subtitle: %s", stream) + if not self._include_srt and stream.extension == "srt": + logger.debug("Ignoring SRT: %s", stream) continue if stream.language not in languages: continue disposition = stream.disposition - if disposition.generic or disposition.hearing_impaired: + + if only_forced and not disposition.forced: + continue + + if ( + disposition.generic + or disposition.hearing_impaired + or (disposition.forced and also_forced) + ): logger.debug("Appending subtitle: %s", stream) subtitles.append(EmbeddedSubtitle(stream, video, {"hash"})) else: diff --git a/tests/subliminal_patch/test_embeddedsubtitles.py b/tests/subliminal_patch/test_embeddedsubtitles.py index 59ffaf0f2..ee1c8d6eb 100644 --- a/tests/subliminal_patch/test_embeddedsubtitles.py +++ b/tests/subliminal_patch/test_embeddedsubtitles.py @@ -36,6 +36,17 @@ def video_multiple_languages(): ) +@pytest.fixture +def config(tmpdir): + return { + "include_ass": True, + "include_srt": True, + "cache_dir": tmpdir, + "ffprobe_path": None, + "ffmpeg_path": None, + } + + @pytest.fixture def video_inexistent(tmpdir): return Movie( @@ -46,12 +57,33 @@ def video_inexistent(tmpdir): ) +def test_init(config): + with EmbeddedSubtitlesProvider(**config) as provider: + assert provider is not None + + def test_inexistent_video(video_inexistent): with EmbeddedSubtitlesProvider() as provider: subtitles = provider.list_subtitles(video_inexistent, {}) assert len(subtitles) == 0 +def test_list_subtitles_only_forced(video_single_language): + with EmbeddedSubtitlesProvider() as provider: + language = Language.fromalpha2("en") + language = Language.rebuild(language, forced=True) + subs = provider.list_subtitles(video_single_language, {language}) + assert len(subs) == 0 + + +def test_list_subtitles_also_forced(video_single_language): + with EmbeddedSubtitlesProvider() as provider: + language_1 = Language.fromalpha2("en") + language_2 = Language.rebuild(language_1, forced=True) + subs = provider.list_subtitles(video_single_language, {language_1, language_2}) + assert any(language_1 == sub.language for sub in subs) + + def test_list_subtitles_single_language(video_single_language): with EmbeddedSubtitlesProvider() as provider: subs = provider.list_subtitles( @@ -81,6 +113,14 @@ def test_list_subtitles_wo_ass(video_single_language): assert not subs +def test_list_subtitles_wo_srt(video_multiple_languages): + with EmbeddedSubtitlesProvider(include_srt=False) as provider: + subs = provider.list_subtitles( + video_multiple_languages, {Language.fromalpha2("en")} + ) + assert not subs + + def test_download_subtitle_multiple(video_multiple_languages): with EmbeddedSubtitlesProvider() as provider: languages = {Language.fromalpha2(code) for code in ("en", "it", "fr")} | {