Start subliminal_patch's core refactor

libs/subliminal_patch/core.py

@@ -1,6 +1,5 @@
 # coding=utf-8
 from __future__ import absolute_import
-import codecs
 import json
 import re
 import os
@@ -13,7 +12,6 @@ import operator
 import unicodedata
 
 import itertools
-from six.moves.http_client import ResponseNotReady
 
 import rarfile
 import requests
@@ -27,13 +25,11 @@ from concurrent.futures import as_completed
 
 from .extensions import provider_registry
 from .exceptions import MustGetBlacklisted
-from subliminal.exceptions import ServiceUnavailable, DownloadLimitExceeded
 from subliminal.score import compute_score as default_compute_score
 from subliminal.utils import hash_napiprojekt, hash_opensubtitles, hash_shooter, hash_thesubdb
 from subliminal.video import VIDEO_EXTENSIONS, Video, Episode, Movie
 from subliminal.core import guessit, ProviderPool, io, is_windows_special_path, \
     ThreadPoolExecutor, check_video
-from subliminal_patch.exceptions import TooManyRequests, APIThrottled
 
 from subzero.language import Language, ENDSWITH_LANGUAGECODE_RE, FULL_LANGUAGE_LIST
 try:
@@ -142,8 +138,7 @@ class _Blacklist(list):
 
 
 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
         self.providers = set(providers or [])
 
@@ -160,10 +155,6 @@ class SZProviderPool(ProviderPool):
 
         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()
 
         if not self.throttle_callback:
@@ -266,9 +257,6 @@ class SZProviderPool(ProviderPool):
         :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
 
         # check video validity
@@ -377,13 +365,7 @@ class SZProviderPool(ProviderPool):
         while True:
             tries += 1
             try:
-                if self.pre_download_hook:
-                    self.pre_download_hook(subtitle)
-
                 self[subtitle.provider_name].download_subtitle(subtitle)
-                if self.post_download_hook:
-                    self.post_download_hook(subtitle)
-
                 break
             except (requests.ConnectionError,
                     requests.exceptions.ProxyError,

tests/subliminal_patch/conftest.py

@@ -17,7 +17,7 @@ def vcr_cassette_dir(request):
     return os.path.join("tests/subliminal_patch/cassettes", request.module.__name__)
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def movies():
     return {
         "dune": Movie(
@@ -85,7 +85,7 @@ def movies():
     }
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def episodes():
     return {
         "got_s03e10": Episode(

tests/subliminal_patch/data/dummy.srt

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

tests/subliminal_patch/test_core.py

@@ -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"]: []
+    }