From b90dab03e89e03b73bfbe15741c8108fc0d2a806 Mon Sep 17 00:00:00 2001 From: Vitiko Date: Sat, 1 Jan 2022 17:59:11 -0400 Subject: [PATCH] Add hi_fallback option to Embedded Subtitles Provider --- bazarr/config.py | 1 + bazarr/get_providers.py | 1 + frontend/src/Settings/Providers/list.ts | 3 + .../providers/embeddedsubtitles.py | 30 +++++++- .../test_embeddedsubtitles.py | 71 +++++++++++++++++++ 5 files changed, 103 insertions(+), 3 deletions(-) diff --git a/bazarr/config.py b/bazarr/config.py index 62eba3088..e21161c2d 100644 --- a/bazarr/config.py +++ b/bazarr/config.py @@ -197,6 +197,7 @@ defaults = { 'embeddedsubtitles': { 'include_ass': 'True', 'include_srt': 'True', + 'hi_fallback': 'False' }, 'subsync': { 'use_subsync': 'False', diff --git a/bazarr/get_providers.py b/bazarr/get_providers.py index fd21cd380..a7c99dbe2 100644 --- a/bazarr/get_providers.py +++ b/bazarr/get_providers.py @@ -202,6 +202,7 @@ def get_providers_auth(): 'embeddedsubtitles': { 'include_ass': settings.embeddedsubtitles.getboolean('include_ass'), 'include_srt': settings.embeddedsubtitles.getboolean('include_srt'), + 'hi_fallback': settings.embeddedsubtitles.getboolean('hi_fallback'), '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 e1e24a96c..d6a2ff2c1 100644 --- a/frontend/src/Settings/Providers/list.ts +++ b/frontend/src/Settings/Providers/list.ts @@ -57,12 +57,15 @@ export const ProviderList: Readonly = [ defaultKey: { include_srt: true, include_ass: true, + hi_fallback: false, }, message: "Warning for cloud users: this provider needs to read the entire file in order to extract subtitles.", keyNameOverride: { include_srt: "Include SRT", include_ass: "Include ASS (will be converted to SRT)", + hi_fallback: + "Use HI subtitles as a fallback (don't enable it if you have a HI language profile)", }, }, { diff --git a/libs/subliminal_patch/providers/embeddedsubtitles.py b/libs/subliminal_patch/providers/embeddedsubtitles.py index d634eb735..4f3a5dfb5 100644 --- a/libs/subliminal_patch/providers/embeddedsubtitles.py +++ b/libs/subliminal_patch/providers/embeddedsubtitles.py @@ -71,12 +71,14 @@ class EmbeddedSubtitlesProvider(Provider): cache_dir=None, ffprobe_path=None, ffmpeg_path=None, + hi_fallback=False, ): 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() ) + self._hi_fallback = hi_fallback self._cached_paths = {} fese.FFPROBE_PATH = ffprobe_path or fese.FFPROBE_PATH @@ -110,7 +112,7 @@ class EmbeddedSubtitlesProvider(Provider): only_forced = all(lang.forced for lang in languages) also_forced = any(lang.forced for lang in languages) - subtitles = [] + allowed_streams = [] for stream in streams: if not self._include_ass and stream.extension == "ass": @@ -135,11 +137,14 @@ class EmbeddedSubtitlesProvider(Provider): or (disposition.forced and also_forced) ): logger.debug("Appending subtitle: %s", stream) - subtitles.append(EmbeddedSubtitle(stream, video, {"hash"})) + allowed_streams.append(stream) else: logger.debug("Ignoring unwanted subtitle: %s", stream) - return subtitles + if self._hi_fallback: + _check_hi_fallback(allowed_streams, languages) + + return [EmbeddedSubtitle(stream, video, {"hash"}) for stream in allowed_streams] def list_subtitles(self, video, languages): if not os.path.isfile(video.original_path): @@ -182,3 +187,22 @@ class EmbeddedSubtitlesProvider(Provider): def _check_allowed_extensions(subtitle: FFprobeSubtitleStream): return subtitle.extension in ("ass", "srt") + + +def _check_hi_fallback(streams, languages): + for language in languages: + compatible_streams = [ + stream for stream in streams if stream.language == language + ] + if len(compatible_streams) == 1: + stream = compatible_streams[0] + logger.debug("HI fallback: updating %s HI to False", stream) + stream.disposition.hearing_impaired = False + + elif all(stream.disposition.hearing_impaired for stream in streams): + for stream in streams: + logger.debug("HI fallback: updating %s HI to False", stream) + stream.disposition.hearing_impaired = False + + else: + logger.debug("HI fallback not needed: %s", compatible_streams) diff --git a/tests/subliminal_patch/test_embeddedsubtitles.py b/tests/subliminal_patch/test_embeddedsubtitles.py index 6c546a61f..c9c8e7cd5 100644 --- a/tests/subliminal_patch/test_embeddedsubtitles.py +++ b/tests/subliminal_patch/test_embeddedsubtitles.py @@ -3,6 +3,8 @@ import os import fese import pytest +import fese +from fese import FFprobeSubtitleStream from subliminal_patch.core import Episode, Movie from subliminal_patch.providers.embeddedsubtitles import EmbeddedSubtitlesProvider from subzero.language import Language @@ -44,6 +46,7 @@ def config(tmpdir): "cache_dir": tmpdir, "ffprobe_path": None, "ffmpeg_path": None, + "hi_fallback": False, } @@ -68,6 +71,74 @@ def test_inexistent_video(video_inexistent): assert len(subtitles) == 0 +@pytest.fixture +def fake_streams(): + return { + "en_hi": FFprobeSubtitleStream( + { + "index": 3, + "codec_name": "subrip", + "disposition": {"default": 1, "hearing_impaired": 1}, + "tags": {"language": "eng", "title": "English"}, + } + ), + "en": FFprobeSubtitleStream( + { + "index": 3, + "codec_name": "subrip", + "tags": {"language": "eng", "title": "English"}, + } + ), + } + + +def test_list_subtitles_hi_fallback_one_stream( + video_single_language, fake_streams, mocker +): + with EmbeddedSubtitlesProvider(hi_fallback=True) as provider: + language = Language.fromalpha2("en") + mocker.patch( + "fese.FFprobeVideoContainer.get_subtitles", + return_value=[fake_streams["en_hi"]], + ) + fake = fese.FFprobeVideoContainer.get_subtitles("")[0] + assert fake.disposition.hearing_impaired == True + + subs = provider.list_subtitles(video_single_language, {language}) + assert subs + assert subs[0].hearing_impaired == False + + +def test_list_subtitles_hi_fallback_multiple_streams( + video_single_language, fake_streams, mocker +): + with EmbeddedSubtitlesProvider(hi_fallback=True) as provider: + language = Language.fromalpha2("en") + mocker.patch( + "fese.FFprobeVideoContainer.get_subtitles", + return_value=[fake_streams["en_hi"], fake_streams["en"]], + ) + subs = provider.list_subtitles(video_single_language, {language}) + assert len(subs) == 2 + assert subs[0].hearing_impaired == True + assert subs[1].hearing_impaired == False + + +def test_list_subtitles_hi_fallback_multiple_hi_streams( + video_single_language, fake_streams, mocker +): + with EmbeddedSubtitlesProvider(hi_fallback=True) as provider: + language = Language.fromalpha2("en") + mocker.patch( + "fese.FFprobeVideoContainer.get_subtitles", + return_value=[fake_streams["en_hi"], fake_streams["en_hi"]], + ) + subs = provider.list_subtitles(video_single_language, {language}) + assert len(subs) == 2 + assert subs[0].hearing_impaired == False + assert subs[1].hearing_impaired == False + + def test_list_subtitles_only_forced(video_single_language): with EmbeddedSubtitlesProvider() as provider: language = Language.fromalpha2("en")