Embedded Subtitles provider: improve HI fallback

This commit is contained in:
Vitiko 2022-05-16 23:56:52 -04:00
parent 3d7fa06f41
commit b96fd93269
2 changed files with 48 additions and 11 deletions

View File

@ -245,21 +245,26 @@ def _check_allowed_extensions(subtitle: FFprobeSubtitleStream):
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
logger.debug("Checking HI fallback for '%s' language", language)
elif all(stream.disposition.hearing_impaired for stream in streams):
for stream in streams:
logger.debug("HI fallback: updating %s HI to False", stream)
streams_ = [stream for stream in streams if stream.language == language]
if len(streams_) == 1 and streams_[0].disposition.hearing_impaired:
logger.debug(
"HI fallback: updating %s HI to False (only subtitle found is HI)",
streams_[0],
)
streams_[0].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 (all subtitles are HI)",
stream,
)
stream.disposition.hearing_impaired = False
else:
logger.debug("HI fallback not needed: %s", compatible_streams)
logger.debug("HI fallback not needed: %s", streams_)
def _discard_possible_incomplete_subtitles(streams):

View File

@ -94,6 +94,21 @@ def fake_streams():
"tags": {"language": "eng", "title": "English"},
}
),
"es_hi": FFprobeSubtitleStream(
{
"index": 3,
"codec_name": "subrip",
"disposition": {"default": 1, "hearing_impaired": 1},
"tags": {"language": "spa", "title": "Spanish"},
}
),
"es": FFprobeSubtitleStream(
{
"index": 3,
"codec_name": "subrip",
"tags": {"language": "spa", "title": "Spanish"},
}
),
}
@ -131,6 +146,23 @@ def test_list_subtitles_hi_fallback_multiple_streams(
assert subs[1].hearing_impaired == False
def test_list_subtitles_hi_fallback_multiple_language_streams(
video_single_language, fake_streams, mocker
):
with EmbeddedSubtitlesProvider(hi_fallback=True) as provider:
languages = {Language.fromalpha2("en"), Language.fromalpha2("es")}
mocker.patch(
# "fese.FFprobeVideoContainer.get_subtitles",
"subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles",
return_value=[fake_streams["en_hi"], fake_streams["es"], fake_streams["es_hi"]],
)
subs = provider.list_subtitles(video_single_language, languages)
assert len(subs) == 3
assert subs[0].hearing_impaired == False # English subittle
assert subs[1].hearing_impaired == False # Spanish subtitle
assert subs[2].hearing_impaired == True # Spanish HI subtitle
def test_list_subtitles_hi_fallback_multiple_hi_streams(
video_single_language, fake_streams, mocker
):