Argenteam provider: fix decoding error

pull/1632/head
Vitiko 3 years ago
parent e1386aedc0
commit d2bbc479bc

@ -7,8 +7,7 @@ import io
import time import time
import urllib.parse import urllib.parse
from json.decoder import JSONDecodeError from simplejson.errors import JSONDecodeError
from zipfile import ZipFile from zipfile import ZipFile
from guessit import guessit from guessit import guessit
from requests import Session from requests import Session
@ -19,8 +18,8 @@ from subliminal_patch.subtitle import Subtitle, guess_matches
from subliminal_patch.providers.mixins import ProviderSubtitleArchiveMixin from subliminal_patch.providers.mixins import ProviderSubtitleArchiveMixin
from subzero.language import Language from subzero.language import Language
BASE_URL = "https://argenteam.net/" BASE_URL = "https://argenteam.net"
API_URL = BASE_URL + "api/v1/" API_URL = f"{BASE_URL}/api/v1"
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -69,10 +68,9 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
multi_result_throttle = 2 # seconds multi_result_throttle = 2 # seconds
def __init__(self): def __init__(self):
self.session = None self.session = Session()
def initialize(self): def initialize(self):
self.session = Session()
self.session.headers.update( self.session.headers.update(
{"User-Agent": os.environ.get("SZ_USER_AGENT", "Sub-Zero/2")} {"User-Agent": os.environ.get("SZ_USER_AGENT", "Sub-Zero/2")}
) )
@ -80,7 +78,105 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
def terminate(self): def terminate(self):
self.session.close() self.session.close()
def search_ids(self, title, **kwargs): def query(self, title, video, titles=None):
is_episode = isinstance(video, Episode)
season = episode = None
url = f"{API_URL}/movie"
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
)
else:
argenteam_ids = self._search_ids(
title, year=video.year, imdb_id=video.imdb_id, titles=titles
)
if not argenteam_ids:
return []
language = self.language_list[0]
subtitles = []
has_multiple_ids = len(argenteam_ids) > 1
for aid in argenteam_ids:
response = self.session.get(url, params={"id": aid}, timeout=10)
response.raise_for_status()
try:
content = response.json()
except JSONDecodeError:
continue
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"),
)
subtitles.append(
ArgenteamSubtitle(
language,
page_link,
download_link,
release_info,
matches_,
)
)
if has_multiple_ids:
time.sleep(self.multi_result_throttle)
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 []
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)
def _search_ids(self, title, **kwargs):
query = title query = title
titles = kwargs.get("titles") or [] titles = kwargs.get("titles") or []
@ -91,7 +187,7 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
logger.debug(f"Searching ID (episode: {is_episode}) for {query}") logger.debug(f"Searching ID (episode: {is_episode}) for {query}")
r = self.session.get(API_URL + "search", params={"q": query}, timeout=10) r = self.session.get(f"{API_URL}/search", params={"q": query}, timeout=10)
r.raise_for_status() r.raise_for_status()
try: try:
@ -131,7 +227,7 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
return match_ids return match_ids
def get_query_matches(self, video, **kwargs): def _get_query_matches(self, video, **kwargs):
matches = set() matches = set()
if isinstance(video, Episode) and kwargs.get("movie_kind") == "episode": if isinstance(video, Episode) and kwargs.get("movie_kind") == "episode":
if video.series and ( if video.series and (
@ -171,107 +267,9 @@ class ArgenteamProvider(Provider, ProviderSubtitleArchiveMixin):
return matches return matches
def combine_release_info(self, release_dict): def _combine_release_info(self, release_dict):
keys = ("source", "codec", "tags", "team") keys = ("source", "codec", "tags", "team")
combine = [release_dict.get(key) for key in keys if release_dict.get(key)] combine = [release_dict.get(key) for key in keys if release_dict.get(key)]
if combine: if combine:
return ".".join(combine) return ".".join(combine)
return "Unknown" return "Unknown"
def query(self, title, video, titles=None):
is_episode = isinstance(video, Episode)
season = episode = None
url = API_URL + "movie"
if is_episode:
season = video.season
episode = video.episode
url = API_URL + "episode"
argenteam_ids = self.search_ids(
title, season=season, episode=episode, titles=titles
)
else:
argenteam_ids = self.search_ids(
title, year=video.year, imdb_id=video.imdb_id, titles=titles
)
if not argenteam_ids:
return []
language = self.language_list[0]
subtitles = []
has_multiple_ids = len(argenteam_ids) > 1
for aid in argenteam_ids:
response = self.session.get(url, params={"id": aid}, timeout=10)
response.raise_for_status()
try:
content = response.json()
except JSONDecodeError:
continue
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"),
)
subtitles.append(
ArgenteamSubtitle(
language,
page_link,
download_link,
release_info,
matches_,
)
)
if has_multiple_ids:
time.sleep(self.multi_result_throttle)
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 []
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)

@ -44,7 +44,7 @@ ftfy_defaults = {
class Subtitle(Subtitle_): class Subtitle(Subtitle_):
storage_path = None storage_path = None
release_info = None release_info = None
matches = None matches = {}
hash_verifiable = False hash_verifiable = False
hearing_impaired_verifiable = False hearing_impaired_verifiable = False
mods = None mods = None

Loading…
Cancel
Save