Add Embedded Subtitles provider memoization/cache

pull/1695/head
Vitiko 3 years ago
parent ddacff3bdc
commit 111998eb79

@ -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)
@functools.lru_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")

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

Loading…
Cancel
Save