From b96fd9326979b08876ef9818bfab0a735dcc0f06 Mon Sep 17 00:00:00 2001
From: Vitiko <averroista@protonmail.com>
Date: Mon, 16 May 2022 23:56:52 -0400
Subject: [PATCH] Embedded Subtitles provider: improve HI fallback

---
 .../providers/embeddedsubtitles.py            | 29 ++++++++++-------
 .../test_embeddedsubtitles.py                 | 32 +++++++++++++++++++
 2 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/libs/subliminal_patch/providers/embeddedsubtitles.py b/libs/subliminal_patch/providers/embeddedsubtitles.py
index 63e7c8fde..d738ee024 100644
--- a/libs/subliminal_patch/providers/embeddedsubtitles.py
+++ b/libs/subliminal_patch/providers/embeddedsubtitles.py
@@ -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
-
-        elif all(stream.disposition.hearing_impaired for stream in streams):
-            for stream in streams:
-                logger.debug("HI fallback: updating %s HI to False", stream)
+        logger.debug("Checking HI fallback for '%s' language", language)
+
+        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):
diff --git a/tests/subliminal_patch/test_embeddedsubtitles.py b/tests/subliminal_patch/test_embeddedsubtitles.py
index 0bda12072..998964e92 100644
--- a/tests/subliminal_patch/test_embeddedsubtitles.py
+++ b/tests/subliminal_patch/test_embeddedsubtitles.py
@@ -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
 ):