Add hi_fallback option to Embedded Subtitles Provider

pull/1658/head
Vitiko 3 years ago
parent e239fc2826
commit b90dab03e8

@ -197,6 +197,7 @@ defaults = {
'embeddedsubtitles': {
'include_ass': 'True',
'include_srt': 'True',
'hi_fallback': 'False'
},
'subsync': {
'use_subsync': 'False',

@ -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"),

@ -57,12 +57,15 @@ export const ProviderList: Readonly<ProviderInfo[]> = [
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)",
},
},
{

@ -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)

@ -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")

Loading…
Cancel
Save