open-source-mirrors
/
bazarr
mirror of https://github.com/morpheus65535/bazarr
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor Argenteam Provider

Browse Source 
* Deprecate text search in favour of IMDB search
* Simplify code
pull/1991/head
Vitiko 2 years ago
parent 2e4480dd5f
commit 52760d8bc7
3 changed files with 142 additions and 199 deletions
Show Stats Download Patch File Download Diff File

256
libs/subliminal_patch/providers/argenteam.py
Unescape View File

@ -1,21 +1,20 @@
# coding=utf-8
from __future__ import absolute_import
from json import JSONDecodeError
import logging
import os
import io
import time
import urllib.parse
from json import JSONDecodeError
from zipfile import ZipFile
from guessit import guessit
from requests import Session
from subliminal import Episode, Movie
from subliminal.utils import sanitize
from subliminal import Episode
from subliminal import Movie
from subliminal_patch.providers import Provider
from subliminal_patch.subtitle import Subtitle, guess_matches
from subliminal_patch.providers.mixins import ProviderSubtitleArchiveMixin
from subliminal_patch.providers.utils import get_archive_from_bytes
from subliminal_patch.providers.utils import get_subtitle_from_archive
from subliminal_patch.providers.utils import update_matches
from subliminal_patch.subtitle import Subtitle
from subzero.language import Language
BASE_URL = "https://argenteam.net"
@ -30,42 +29,31 @@ class ArgenteamSubtitle(Subtitle):
def __init__(self, language, page_link, download_link, release_info, matches):
super(ArgenteamSubtitle, self).__init__(language, page_link=page_link)
self._found_matches = matches
self.page_link = page_link
self.download_link = download_link
self.found_matches = matches
self._release_info = release_info
# Original subtitle filename guessed from the URL
self.release_info = urllib.parse.unquote(self.download_link.split("/")[-1])
self.release_info = release_info
@property
def id(self):
return self.download_link
def get_matches(self, video):
type_ = "episode" if isinstance(video, Episode) else "movie"
update_matches(self._found_matches, video, self.release_info)
self.found_matches |= guess_matches(
video,
guessit(self.release_info, {"type": type_}),
)
self.found_matches |= guess_matches(
video,
guessit(self._release_info, {"type": type_}),
)
return self.found_matches
return self._found_matches
class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
provider_name = "argenteam"
# Safe to assume every subtitle from Argenteam is Latam Spanish
languages = {Language("spa", "MX")}
video_types = (Episode, Movie)
subtitle_class = ArgenteamSubtitle
hearing_impaired_verifiable = False
language_list = list(languages)
multi_result_throttle = 2 # seconds
_default_lang = Language("spa", "MX")
def __init__(self):
self.session = Session()
@ -78,31 +66,36 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
def terminate(self):
self.session.close()
def query(self, title, video, titles=None):
def query(self, video):
is_episode = isinstance(video, Episode)
season = episode = None
url = f"{API_URL}/movie"
imdb_id = video.series_imdb_id if is_episode else video.imdb_id
if not imdb_id:
logger.debug("%s doesn't have IMDB ID. Can't search")
return []
if is_episode:
season = video.season
episode = video.episode
url = f"{API_URL}/episode"
argenteam_ids = self._search_ids(
title, season=season, episode=episode, titles=titles
imdb_id, season=video.season, episode=video.episode
)
else:
argenteam_ids = self._search_ids(
title, year=video.year, imdb_id=video.imdb_id, titles=titles
)
argenteam_ids = self._search_ids(imdb_id)
if not argenteam_ids:
logger.debug("No IDs found")
return []
language = self.language_list[0]
return self._parse_subtitles(argenteam_ids, is_episode)
def _parse_subtitles(self, ids, is_episode=True):
movie_kind = "episode" if is_episode else "movie"
subtitles = []
has_multiple_ids = len(argenteam_ids) > 1
for aid in argenteam_ids:
response = self.session.get(url, params={"id": aid}, timeout=10)
for aid in ids:
response = self.session.get(
f"{API_URL}/{movie_kind}", params={"id": aid}, timeout=10
)
response.raise_for_status()
try:
@ -113,81 +106,55 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
if not content or not content.get("releases"):
continue
imdb_id = year = None
returned_title = title
if not is_episode and "info" in content:
imdb_id = content["info"].get("imdb")
year = content["info"].get("year")
returned_title = content["info"].get("title", title)
for r in content["releases"]:
for s in r["subtitles"]:
movie_kind = "episode" if is_episode else "movie"
page_link = f"{BASE_URL}/{movie_kind}/{aid}"
release_info = self._combine_release_info(r)
download_link = s["uri"].replace("http://", "https://")
matches_ = self._get_query_matches(
video,
movie_kind=movie_kind,
season=season,
episode=episode,
title=returned_title,
year=year,
imdb_id=imdb_id,
tvdb_id=content.get("tvdb"),
)
release_info = self._combine_release_info(r, s)
if matches_ is not None:
subtitles.append(
ArgenteamSubtitle(
language,
page_link,
download_link,
release_info,
matches_,
)
)
logger.debug("Got release info: %s", release_info)
if has_multiple_ids:
time.sleep(self.multi_result_throttle)
download_link = s["uri"].replace("http://", "https://")
# Already matched within query
if is_episode:
matches = {"series", "title", "season", "episode", "imdb_id"}
else:
matches = {"title", "year", "imdb_id"}
subtitles.append(
ArgenteamSubtitle(
self._default_lang,
page_link,
download_link,
release_info,
matches,
)
)
return subtitles
def list_subtitles(self, video, languages):
if isinstance(video, Episode):
titles = [video.series] + video.alternative_series[:2]
else:
titles = [video.title] + video.alternative_titles[:2]
for title in titles:
subs = self.query(title, video, titles=titles)
if subs:
return subs
time.sleep(self.multi_result_throttle)
return []
return self.query(video)
def download_subtitle(self, subtitle):
# download as a zip
logger.info("Downloading subtitle %r", subtitle)
r = self.session.get(subtitle.download_link, timeout=10)
r.raise_for_status()
# open the zip
with ZipFile(io.BytesIO(r.content)) as zf:
subtitle.content = self.get_subtitle_from_archive(subtitle, zf)
archive = get_archive_from_bytes(r.content)
subtitle.content = get_subtitle_from_archive(archive)
def _search_ids(self, title, **kwargs):
query = title
titles = kwargs.get("titles") or []
def _search_ids(self, identifier, **kwargs):
"""
:param identifier: imdb_id or title (without year)
"""
identifier = identifier.lstrip("tt")
is_episode = False
query = identifier
if kwargs.get("season") and kwargs.get("episode"):
is_episode = True
query = f"{title} S{kwargs['season']:02}E{kwargs['episode']:02}"
query = f"{identifier} S{kwargs['season']:02}E{kwargs['episode']:02}"
logger.debug(f"Searching ID (episode: {is_episode}) for {query}")
logger.debug("Searching ID for %s", query)
r = self.session.get(f"{API_URL}/search", params={"q": query}, timeout=10)
r.raise_for_status()
@ -200,84 +167,27 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
if not results.get("results"):
return []
match_ids = []
for result in results["results"]:
if result["type"] == "movie" and is_episode:
continue
imdb = f"tt{result.get('imdb', 'n/a')}"
if not is_episode and imdb == kwargs.get("imdb_id"):
logger.debug("Movie matched by IMDB ID, taking shortcut")
match_ids = [result["id"]]
break
# advanced title check in case of multiple movie results
title_year = kwargs.get("year") and kwargs.get("title")
if results["total"] > 1 and not is_episode and title_year:
sanitized = sanitize(result["title"])
titles = [f"{sanitize(name)} {kwargs['year']}" for name in titles]
if sanitized not in titles:
continue
match_ids.append(result["id"])
if match_ids:
ids = ", ".join(str(id) for id in match_ids)
logger.debug("Found matching IDs: %s", ids)
else:
logger.debug("Nothing found from %s query", query)
match_ids = [result["id"] for result in results["results"]]
logger.debug("Found matching IDs: %s", match_ids)
return match_ids
def _get_query_matches(self, video, **kwargs):
matches = set()
if isinstance(video, Episode) and kwargs.get("movie_kind") == "episode":
if (kwargs.get("tvdb_id") and video.series_tvdb_id) and str(
video.series_tvdb_id
) != str(kwargs.get("tvdb_id")):
logger.debug(
"TVDB ID not matched: %s - %s", kwargs, video.series_tvdb_id
)
return None
if video.series and (
sanitize(kwargs.get("title"))
in (
sanitize(name) for name in [video.series] + video.alternative_series
)
):
matches.add("series")
if video.season and kwargs.get("season") == video.season:
matches.add("season")
if video.episode and kwargs.get("episode") == video.episode:
matches.add("episode")
# year (year is not available for series, but we assume it matches)
matches.add("year")
elif isinstance(video, Movie) and kwargs.get("movie_kind") == "movie":
if video.title and (
sanitize(kwargs.get("title"))
in (sanitize(name) for name in [video.title] + video.alternative_titles)
):
matches.add("title")
if video.imdb_id and f"tt{kwargs.get('imdb_id')}" == str(video.imdb_id):
matches.add("imdb_id")
if video.year and kwargs.get("year") == video.year:
matches.add("year")
else:
logger.info(f"{kwargs.get('movie_kind')} is not a valid movie_kind")
def _combine_release_info(self, release_dict, subtitle_dict):
releases = [
urllib.parse.unquote(subtitle_dict.get("uri", "Unknown").split("/")[-1])
]
return matches
combine = [
release_dict.get(key)
for key in ("source", "codec", "tags")
if release_dict.get(key)
]
def _combine_release_info(self, release_dict):
keys = ("source", "codec", "tags", "team")
combine = [release_dict.get(key) for key in keys if release_dict.get(key)]
if combine:
return ".".join(combine)
return "Unknown"
r_info = ".".join(combine)
if release_dict.get("team"):
r_info += f"-{release_dict['team']}"
releases.append(r_info)
return "\n".join(releases)

1
tests/subliminal_patch/conftest.py
Unescape View File

@ -123,6 +123,7 @@ def episodes():
1,
1,
source="Blu-Ray",
series_imdb_id="tt0903747",
release_group="REWARD",
resolution="720p",
video_codec="H.264",

84
tests/subliminal_patch/test_argenteam.py
Unescape View File

@ -8,14 +8,39 @@ from subliminal_patch.core import Episode
from subzero.language import Language
@pytest.mark.parametrize(
"imdb_id,expected_id", [("tt0028950", 62790), ("tt0054407", 102006)]
)
def test_search_ids_movie(imdb_id, expected_id):
with ArgenteamProvider() as provider:
ids = provider._search_ids(imdb_id)
assert ids[0] == expected_id
def test_search_ids_tv_show():
with ArgenteamProvider() as provider:
ids = provider._search_ids("tt0306414", season=1, episode=1)
assert ids[0] == 10075
def test_parse_subtitles_episode():
with ArgenteamProvider() as provider:
assert len(provider._parse_subtitles([10075])) > 1
def test_parse_subtitles_movie():
with ArgenteamProvider() as provider:
assert len(provider._parse_subtitles([61], is_episode=False)) > 3
def test_get_matches_episode(episodes):
episode = episodes["breaking_bad_s01e01"]
subtitle = ArgenteamSubtitle(
Language.fromalpha2("es"),
None,
"https://argenteam.net/subtitles/24002/Breaking.Bad.%282008%29.S01E01-Pilot.BluRay.x264.720p-REWARD",
"BluRay x264 720p",
{"title", "season", "episode", "imdb_id"},
"Breaking.Bad.(2008).S01E01-Pilot.BluRay.x264.720p-REWARD\nBluRay x264 720p",
{"series", "title", "season", "episode", "imdb_id"},
)
matches = subtitle.get_matches(episode)
assert matches == {
@ -52,10 +77,10 @@ def test_get_matches_movie(movies):
"resolution",
"edition",
"video_codec",
"streaming_service",
}
@pytest.mark.vcr
def test_list_subtitles_movie(movies):
item = movies["dune"]
with ArgenteamProvider() as provider:
@ -69,7 +94,20 @@ def test_list_subtitles_movie(movies):
assert any(expected == sub.download_link for sub in subtitles)
@pytest.mark.vcr
def test_list_subtitles_movie_no_imdb(movies):
item = movies["dune"]
item.imdb_id = None
with ArgenteamProvider() as provider:
assert not provider.list_subtitles(item, {Language("spa", "MX")})
def test_list_subtitles_movie_not_found(movies):
item = movies["dune"]
item.imdb_id = "tt29318321832"
with ArgenteamProvider() as provider:
assert not provider.list_subtitles(item, {Language("spa", "MX")})
def test_list_subtitles_episode(episodes):
item = episodes["breaking_bad_s01e01"]
with ArgenteamProvider() as provider:
@ -82,29 +120,23 @@ def test_list_subtitles_episode(episodes):
assert any(expected == sub.download_link for sub in subtitles)
@pytest.mark.vcr
def test_list_subtitles_episode_no_imdb_id(episodes):
item = episodes["breaking_bad_s01e01"]
item.series_imdb_id = None
with ArgenteamProvider() as provider:
assert not provider.list_subtitles(item, {Language("spa", "MX")})
def test_list_subtitles_episode_not_found(episodes):
item = episodes["breaking_bad_s01e01"]
item.series_imdb_id = "tt29318321832"
with ArgenteamProvider() as provider:
assert not provider.list_subtitles(item, {Language("spa", "MX")})
def test_download_subtitle(episodes):
item = episodes["breaking_bad_s01e01"]
with ArgenteamProvider() as provider:
subtitles = provider.list_subtitles(item, {Language("spa", "MX")})
subtitle = subtitles[0]
provider.download_subtitle(subtitle)
assert subtitle.content is not None
@pytest.mark.vcr
def test_list_subtitles_episode_with_tvdb():
video = Episode(
"Severance.S01E01.720p.BluRay.X264-REWARD.mkv",
"Severance",
1,
1,
source="Blu-Ray",
release_group="REWARD",
resolution="720p",
video_codec="H.264",
series_tvdb_id=371980,
)
with ArgenteamProvider() as provider:
subtitles = provider.list_subtitles(video, {Language("spa", "MX")})
assert len(subtitles) == 0
provider.download_subtitle(subtitles[0])
assert subtitles[0].is_valid()

Write Preview
Loading…
Cancel
Save