Start subliminal_patch's core refactor

subliminal_patch_core_refactor
vitiko98 2 years ago
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,9 +257,6 @@ 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 = self.language_hook(provider)
else:
languages_search_base = languages languages_search_base = languages
# check video validity # check video validity
@ -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…
Cancel
Save