mirror of https://github.com/morpheus65535/bazarr
Start subliminal_patch's core refactor
This commit is contained in:
parent
2bd38817b0
commit
949fd26b24
|
@ -1,6 +1,5 @@
|
||||||
# coding=utf-8
|
# coding=utf-8
|
||||||
from __future__ import absolute_import
|
from __future__ import absolute_import
|
||||||
import codecs
|
|
||||||
import json
|
import json
|
||||||
import re
|
import re
|
||||||
import os
|
import os
|
||||||
|
@ -13,7 +12,6 @@ import operator
|
||||||
import unicodedata
|
import unicodedata
|
||||||
|
|
||||||
import itertools
|
import itertools
|
||||||
from six.moves.http_client import ResponseNotReady
|
|
||||||
|
|
||||||
import rarfile
|
import rarfile
|
||||||
import requests
|
import requests
|
||||||
|
@ -27,13 +25,11 @@ from concurrent.futures import as_completed
|
||||||
|
|
||||||
from .extensions import provider_registry
|
from .extensions import provider_registry
|
||||||
from .exceptions import MustGetBlacklisted
|
from .exceptions import MustGetBlacklisted
|
||||||
from subliminal.exceptions import ServiceUnavailable, DownloadLimitExceeded
|
|
||||||
from subliminal.score import compute_score as default_compute_score
|
from subliminal.score import compute_score as default_compute_score
|
||||||
from subliminal.utils import hash_napiprojekt, hash_opensubtitles, hash_shooter, hash_thesubdb
|
from subliminal.utils import hash_napiprojekt, hash_opensubtitles, hash_shooter, hash_thesubdb
|
||||||
from subliminal.video import VIDEO_EXTENSIONS, Video, Episode, Movie
|
from subliminal.video import VIDEO_EXTENSIONS, Video, Episode, Movie
|
||||||
from subliminal.core import guessit, ProviderPool, io, is_windows_special_path, \
|
from subliminal.core import guessit, ProviderPool, io, is_windows_special_path, \
|
||||||
ThreadPoolExecutor, check_video
|
ThreadPoolExecutor, check_video
|
||||||
from subliminal_patch.exceptions import TooManyRequests, APIThrottled
|
|
||||||
|
|
||||||
from subzero.language import Language, ENDSWITH_LANGUAGECODE_RE, FULL_LANGUAGE_LIST
|
from subzero.language import Language, ENDSWITH_LANGUAGECODE_RE, FULL_LANGUAGE_LIST
|
||||||
try:
|
try:
|
||||||
|
@ -142,8 +138,7 @@ class _Blacklist(list):
|
||||||
|
|
||||||
|
|
||||||
class SZProviderPool(ProviderPool):
|
class SZProviderPool(ProviderPool):
|
||||||
def __init__(self, providers=None, provider_configs=None, blacklist=None, ban_list=None, throttle_callback=None,
|
def __init__(self, providers=None, provider_configs=None, blacklist=None, ban_list=None, throttle_callback=None):
|
||||||
pre_download_hook=None, post_download_hook=None, language_hook=None):
|
|
||||||
#: Name of providers to use
|
#: Name of providers to use
|
||||||
self.providers = set(providers or [])
|
self.providers = set(providers or [])
|
||||||
|
|
||||||
|
@ -160,10 +155,6 @@ class SZProviderPool(ProviderPool):
|
||||||
|
|
||||||
self.throttle_callback = throttle_callback
|
self.throttle_callback = throttle_callback
|
||||||
|
|
||||||
self.pre_download_hook = pre_download_hook
|
|
||||||
self.post_download_hook = post_download_hook
|
|
||||||
self.language_hook = language_hook
|
|
||||||
|
|
||||||
self._born = time.time()
|
self._born = time.time()
|
||||||
|
|
||||||
if not self.throttle_callback:
|
if not self.throttle_callback:
|
||||||
|
@ -266,10 +257,7 @@ class SZProviderPool(ProviderPool):
|
||||||
:rtype: list of :class:`~subliminal.subtitle.Subtitle` or None
|
:rtype: list of :class:`~subliminal.subtitle.Subtitle` or None
|
||||||
|
|
||||||
"""
|
"""
|
||||||
if self.language_hook:
|
languages_search_base = languages
|
||||||
languages_search_base = self.language_hook(provider)
|
|
||||||
else:
|
|
||||||
languages_search_base = languages
|
|
||||||
|
|
||||||
# check video validity
|
# check video validity
|
||||||
if not provider_registry[provider].check(video):
|
if not provider_registry[provider].check(video):
|
||||||
|
@ -377,13 +365,7 @@ class SZProviderPool(ProviderPool):
|
||||||
while True:
|
while True:
|
||||||
tries += 1
|
tries += 1
|
||||||
try:
|
try:
|
||||||
if self.pre_download_hook:
|
|
||||||
self.pre_download_hook(subtitle)
|
|
||||||
|
|
||||||
self[subtitle.provider_name].download_subtitle(subtitle)
|
self[subtitle.provider_name].download_subtitle(subtitle)
|
||||||
if self.post_download_hook:
|
|
||||||
self.post_download_hook(subtitle)
|
|
||||||
|
|
||||||
break
|
break
|
||||||
except (requests.ConnectionError,
|
except (requests.ConnectionError,
|
||||||
requests.exceptions.ProxyError,
|
requests.exceptions.ProxyError,
|
||||||
|
|
|
@ -17,7 +17,7 @@ def vcr_cassette_dir(request):
|
||||||
return os.path.join("tests/subliminal_patch/cassettes", request.module.__name__)
|
return os.path.join("tests/subliminal_patch/cassettes", request.module.__name__)
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture(scope="session")
|
||||||
def movies():
|
def movies():
|
||||||
return {
|
return {
|
||||||
"dune": Movie(
|
"dune": Movie(
|
||||||
|
@ -85,7 +85,7 @@ def movies():
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture(scope="session")
|
||||||
def episodes():
|
def episodes():
|
||||||
return {
|
return {
|
||||||
"got_s03e10": Episode(
|
"got_s03e10": Episode(
|
||||||
|
|
|
@ -0,0 +1,40 @@
|
||||||
|
1
|
||||||
|
00:00:00,000 --> 00:00:05,000
|
||||||
|
Foo.
|
||||||
|
|
||||||
|
2
|
||||||
|
00:00:05,000 --> 00:00:07,000
|
||||||
|
Bar.
|
||||||
|
|
||||||
|
3
|
||||||
|
00:00:07,000 --> 00:00:09,000
|
||||||
|
Foo.
|
||||||
|
|
||||||
|
4
|
||||||
|
00:00:09,000 --> 00:00:11,000
|
||||||
|
Bar.
|
||||||
|
|
||||||
|
5
|
||||||
|
00:00:11,000 --> 00:00:13,000
|
||||||
|
Foo.
|
||||||
|
|
||||||
|
6
|
||||||
|
00:00:13,000 --> 00:00:15,000
|
||||||
|
Bar.
|
||||||
|
|
||||||
|
7
|
||||||
|
00:00:15,000 --> 00:00:17,000
|
||||||
|
Foo.
|
||||||
|
|
||||||
|
8
|
||||||
|
00:00:17,000 --> 00:00:19,000
|
||||||
|
Bar.
|
||||||
|
|
||||||
|
9
|
||||||
|
00:00:19,000 --> 00:00:21,000
|
||||||
|
Foo.
|
||||||
|
|
||||||
|
10
|
||||||
|
00:00:21,000 --> 00:00:23,000
|
||||||
|
Bar
|
||||||
|
|
|
@ -0,0 +1,152 @@
|
||||||
|
import os
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from subliminal_patch import core
|
||||||
|
from subliminal_patch.score import compute_score
|
||||||
|
from subliminal_patch import core_persistent as corep
|
||||||
|
from subliminal_patch import Provider
|
||||||
|
from subliminal_patch import Subtitle
|
||||||
|
from subliminal_patch.core import SZProviderPool as Pool
|
||||||
|
from subzero.language import Language
|
||||||
|
|
||||||
|
|
||||||
|
class FakeProviderSubtitle(Subtitle):
|
||||||
|
provider_name = "fake"
|
||||||
|
|
||||||
|
def __init__(self, language, id, matches=None):
|
||||||
|
super().__init__(language, page_link=id)
|
||||||
|
self._id = id
|
||||||
|
self._matches = set(matches or [])
|
||||||
|
|
||||||
|
self.release_info = id
|
||||||
|
|
||||||
|
def get_matches(self, video):
|
||||||
|
return self._matches
|
||||||
|
|
||||||
|
def id(self):
|
||||||
|
return self.id
|
||||||
|
|
||||||
|
|
||||||
|
_ENGLISH = Language.fromietf("en")
|
||||||
|
|
||||||
|
|
||||||
|
class FakeProvider(Provider):
|
||||||
|
languages = {_ENGLISH}
|
||||||
|
video_types = (core.Movie, core.Episode)
|
||||||
|
|
||||||
|
def __init__(self, fake_subtitles=None):
|
||||||
|
self._fake_subtitles = fake_subtitles
|
||||||
|
|
||||||
|
def initialize(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def terminate(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def list_subtitles(self, video, languages):
|
||||||
|
return self._fake_subtitles
|
||||||
|
|
||||||
|
def download_subtitle(self, subtitle):
|
||||||
|
filename = os.path.join(
|
||||||
|
os.path.abspath(os.path.dirname(__file__)), "data", "dummy.srt"
|
||||||
|
)
|
||||||
|
with open(filename, "rb") as f:
|
||||||
|
subtitle.content = f.read()
|
||||||
|
|
||||||
|
|
||||||
|
core.provider_registry.register("fake", FakeProvider)
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_init_default():
|
||||||
|
with Pool() as pool:
|
||||||
|
assert pool
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_init_w_providers():
|
||||||
|
with Pool(
|
||||||
|
providers={"opensubtitles"},
|
||||||
|
provider_configs={"opensubtitles": {"username": "foo", "password": "bar"}},
|
||||||
|
) as pool:
|
||||||
|
assert pool.providers is not None
|
||||||
|
assert isinstance(pool.provider_configs, core._ProviderConfigs)
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_init_w_blacklist():
|
||||||
|
with Pool(blacklist=[("foo", "bar")]) as pool:
|
||||||
|
assert isinstance(pool.blacklist, core._Blacklist)
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_init_w_ban_list():
|
||||||
|
with Pool(ban_list={"must_contain": ["foo"], "must_not_contain": ["bar"]}) as pool:
|
||||||
|
assert isinstance(pool.ban_list, core._Banlist)
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_init_w_throttle_callback():
|
||||||
|
with Pool(throttle_callback=None) as pool:
|
||||||
|
pool.throttle_callback("foo", "bar")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def fake_subtitle():
|
||||||
|
yield FakeProviderSubtitle(_ENGLISH, "foo")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def fake_pool(fake_subtitle):
|
||||||
|
with Pool(
|
||||||
|
providers={"fake"},
|
||||||
|
provider_configs={"fake": {"fake_subtitles": [fake_subtitle]}},
|
||||||
|
) as pool:
|
||||||
|
yield pool
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_list_subtitles_provider(fake_pool, fake_subtitle, movies):
|
||||||
|
result = fake_pool.list_subtitles_provider("fake", movies["dune"], {_ENGLISH})
|
||||||
|
assert fake_subtitle in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_list_subtitles(fake_pool, fake_subtitle, movies):
|
||||||
|
result = fake_pool.list_subtitles(movies["dune"], {_ENGLISH})
|
||||||
|
assert fake_subtitle in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_download_subtitle(fake_pool, fake_subtitle):
|
||||||
|
downloaded = fake_pool.download_subtitle(fake_subtitle)
|
||||||
|
assert downloaded is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_download_best_subtitles(fake_pool, fake_subtitle, movies):
|
||||||
|
result = fake_pool.download_best_subtitles(
|
||||||
|
[fake_subtitle], movies["dune"], [_ENGLISH], compute_score=compute_score
|
||||||
|
)
|
||||||
|
assert result == [fake_subtitle]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="module")
|
||||||
|
def empty_pool():
|
||||||
|
with Pool() as pool:
|
||||||
|
yield pool
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_core_persistent_list_all_subtitles(movies, empty_pool):
|
||||||
|
assert corep.list_all_subtitles([movies["dune"]], {_ENGLISH}, empty_pool) == {
|
||||||
|
movies["dune"]: []
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_core_persistent_list_supported_languages(empty_pool):
|
||||||
|
assert corep.list_supported_languages(empty_pool) == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_core_persistent_list_supported_video_types(empty_pool):
|
||||||
|
assert corep.list_supported_video_types(empty_pool) == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_core_persistent_download_subtitles(empty_pool):
|
||||||
|
corep.download_subtitles([], empty_pool)
|
||||||
|
|
||||||
|
|
||||||
|
def test_pool_core_persistent_download_best_subtitles(movies, empty_pool):
|
||||||
|
assert corep.download_best_subtitles([movies["dune"]], {_ENGLISH}, empty_pool) == {
|
||||||
|
movies["dune"]: []
|
||||||
|
}
|
Loading…
Reference in New Issue