Add Embedded Subtitles provider memoization/cache

2 files changed, 52 insertions, 10 deletions

diff --git a/libs/subliminal_patch/providers/embeddedsubtitles.py b/libs/subliminal_patch/providers/embeddedsubtitles.py
index 4ac49a7b8..497a76ecd 100644
--- a/libs/subliminal_patch/providers/embeddedsubtitles.py
+++ b/libs/subliminal_patch/providers/embeddedsubtitles.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+import functools
 import logging
 import os
 import shutil
@@ -8,9 +9,9 @@ import tempfile
 from babelfish import language_converters
 import fese
 from fese import check_integrity
-from fese import InvalidFile
 from fese import FFprobeSubtitleStream
 from fese import FFprobeVideoContainer
+from fese import InvalidFile
 from fese import to_srt
 from subliminal.subtitle import fix_line_ending
 from subliminal_patch.core import Episode
@@ -67,6 +68,7 @@ class EmbeddedSubtitlesProvider(Provider):
 
     video_types = (Episode, Movie)
     subtitle_class = EmbeddedSubtitle
+    _blacklist = set()
 
     def __init__(
         self,
@@ -102,12 +104,17 @@ class EmbeddedSubtitlesProvider(Provider):
         shutil.rmtree(self._cache_dir, ignore_errors=True)
 
     def query(self, path: str, languages, media_type):
-        video = FFprobeVideoContainer(path)
+        if path in self._blacklist:
+            logger.debug("Ignoring blacklisted path: %s", path)
+            return []
+
+        video = _get_memoized_video_container(path)
 
         try:
             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)
+            self._blacklist.add(path)
             streams = []
 
         if not streams:
@@ -148,6 +155,8 @@ class EmbeddedSubtitlesProvider(Provider):
         if self._hi_fallback:
             _check_hi_fallback(allowed_streams, languages)
 
+        logger.debug("Cache info: %s", _get_memoized_video_container.cache_info())
+
         return [
             EmbeddedSubtitle(stream, video, {"hash"}, media_type)
             for stream in allowed_streams
@@ -199,6 +208,17 @@ class EmbeddedSubtitlesProvider(Provider):
         return new_subtitle_path
 
 
+class _MemoizedFFprobeVideoContainer(FFprobeVideoContainer):
+    @functools.lru_cache
+    def get_subtitles(self, *args, **kwargs):
+        return super().get_subtitles(*args, **kwargs)
+
+
[email protected]_cache(maxsize=8096)
+def _get_memoized_video_container(path: str):
+    return _MemoizedFFprobeVideoContainer(path)
+
+
 def _check_allowed_extensions(subtitle: FFprobeSubtitleStream):
     return subtitle.extension in ("ass", "srt")
 
diff --git a/tests/subliminal_patch/test_embeddedsubtitles.py b/tests/subliminal_patch/test_embeddedsubtitles.py
index e42330425..0301c2f0b 100644
--- a/tests/subliminal_patch/test_embeddedsubtitles.py
+++ b/tests/subliminal_patch/test_embeddedsubtitles.py
@@ -1,13 +1,18 @@
 # -*- coding: utf-8 -*-
 import os
+import tempfile
 
 import fese
-import pytest
-import fese
 from fese import FFprobeSubtitleStream
-from subliminal_patch.core import Episode, Movie
+import pytest
+import subliminal_patch
+from subliminal_patch.core import Episode
+from subliminal_patch.core import Movie
 from subliminal_patch.exceptions import MustGetBlacklisted
-from subliminal_patch.providers.embeddedsubtitles import EmbeddedSubtitlesProvider
+from subliminal_patch.providers.embeddedsubtitles import \
+    _MemoizedFFprobeVideoContainer
+from subliminal_patch.providers.embeddedsubtitles import \
+    EmbeddedSubtitlesProvider
 from subzero.language import Language
 
 _DATA = os.path.join(os.path.abspath(os.path.dirname(__file__)), "data")
@@ -99,10 +104,11 @@ def test_list_subtitles_hi_fallback_one_stream(
     with EmbeddedSubtitlesProvider(hi_fallback=True) as provider:
         language = Language.fromalpha2("en")
         mocker.patch(
-            "fese.FFprobeVideoContainer.get_subtitles",
+            #            "fese.FFprobeVideoContainer.get_subtitles",
+            "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles",
             return_value=[fake_streams["en_hi"]],
         )
-        fake = fese.FFprobeVideoContainer.get_subtitles("")[0]
+        fake = _MemoizedFFprobeVideoContainer.get_subtitles("")[0]
         assert fake.disposition.hearing_impaired == True
 
         subs = provider.list_subtitles(video_single_language, {language})
@@ -116,7 +122,8 @@ def test_list_subtitles_hi_fallback_multiple_streams(
     with EmbeddedSubtitlesProvider(hi_fallback=True) as provider:
         language = Language.fromalpha2("en")
         mocker.patch(
-            "fese.FFprobeVideoContainer.get_subtitles",
+            # "fese.FFprobeVideoContainer.get_subtitles",
+            "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles",
             return_value=[fake_streams["en_hi"], fake_streams["en"]],
         )
         subs = provider.list_subtitles(video_single_language, {language})
@@ -131,7 +138,8 @@ def test_list_subtitles_hi_fallback_multiple_hi_streams(
     with EmbeddedSubtitlesProvider(hi_fallback=True) as provider:
         language = Language.fromalpha2("en")
         mocker.patch(
-            "fese.FFprobeVideoContainer.get_subtitles",
+            # "fese.FFprobeVideoContainer.get_subtitles",
+            "subliminal_patch.providers.embeddedsubtitles._MemoizedFFprobeVideoContainer.get_subtitles",
             return_value=[fake_streams["en_hi"], fake_streams["en_hi"]],
         )
         subs = provider.list_subtitles(video_single_language, {language})
@@ -229,3 +237,17 @@ def test_download_invalid_subtitle(video_single_language):
         except MustGetBlacklisted as error:
             assert error.id == subtitle.id
             assert error.media_type == subtitle.media_type
+
+
+def test_memoized(video_single_language, mocker):
+    with EmbeddedSubtitlesProvider() as provider:
+        provider.list_subtitles(video_single_language, {Language.fromalpha2("en")})
+
+    with EmbeddedSubtitlesProvider() as provider:
+        mocker.patch("fese.FFprobeVideoContainer.get_subtitles")
+        assert (
+            provider.list_subtitles(video_single_language, {Language.fromalpha2("en")})[
+                0
+            ]
+            is not None
+        )