[danbooru] slight code refactoring

This commit is contained in:
Mike Fährmann
2020-12-25 22:06:25 +01:00
parent e41e2be2f9
commit e7d446a8f7

View File

@@ -12,7 +12,6 @@ from .common import Extractor, Message
from .. import text from .. import text
import datetime import datetime
BASE_PATTERN = ( BASE_PATTERN = (
r"(?:https?://)?" r"(?:https?://)?"
r"(danbooru|hijiribe|sonohara|safebooru)" r"(danbooru|hijiribe|sonohara|safebooru)"
@@ -33,7 +32,6 @@ class DanbooruExtractor(Extractor):
super().__init__(match) super().__init__(match)
self.root = "https://{}.donmai.us".format(match.group(1)) self.root = "https://{}.donmai.us".format(match.group(1))
self.ugoira = self.config("ugoira", False) self.ugoira = self.config("ugoira", False)
self.params = {}
username, api_key = self._get_auth_info() username, api_key = self._get_auth_info()
if username: if username:
@@ -71,13 +69,16 @@ class DanbooruExtractor(Extractor):
yield Message.Url, url, post yield Message.Url, url, post
def metadata(self): def metadata(self):
return {} return ()
def posts(self): def posts(self):
return self._pagination(self.root + "/posts.json") return ()
def _pagination(self, url, pagenum=False): def _pagination(self, endpoint, params=None, pagenum=False):
params = self.params.copy() url = self.root + endpoint
if params is None:
params = {}
params["limit"] = self.per_page params["limit"] = self.per_page
params["page"] = self.page_start params["page"] = self.page_start
@@ -122,10 +123,14 @@ class DanbooruTagExtractor(DanbooruExtractor):
def __init__(self, match): def __init__(self, match):
super().__init__(match) super().__init__(match)
self.params["tags"] = text.unquote(match.group(2).replace("+", " ")) self.tags = text.unquote(match.group(2).replace("+", " "))
def metadata(self): def metadata(self):
return {"search_tags": self.params["tags"]} return {"search_tags": self.tags}
def posts(self):
params = {"tags": self.tags}
return self._pagination("/posts.json", params)
class DanbooruPoolExtractor(DanbooruExtractor): class DanbooruPoolExtractor(DanbooruExtractor):
@@ -141,7 +146,6 @@ class DanbooruPoolExtractor(DanbooruExtractor):
def __init__(self, match): def __init__(self, match):
super().__init__(match) super().__init__(match)
self.pool_id = match.group(2) self.pool_id = match.group(2)
self.params["tags"] = "pool:" + self.pool_id
def metadata(self): def metadata(self):
url = "{}/pools/{}.json".format(self.root, self.pool_id) url = "{}/pools/{}.json".format(self.root, self.pool_id)
@@ -150,6 +154,10 @@ class DanbooruPoolExtractor(DanbooruExtractor):
del pool["post_ids"] del pool["post_ids"]
return {"pool": pool} return {"pool": pool}
def posts(self):
params = {"tags": "pool:" + self.pool_id}
return self._pagination("/posts.json", params)
class DanbooruPostExtractor(DanbooruExtractor): class DanbooruPostExtractor(DanbooruExtractor):
"""Extractor for single danbooru posts""" """Extractor for single danbooru posts"""
@@ -193,10 +201,9 @@ class DanbooruPopularExtractor(DanbooruExtractor):
def __init__(self, match): def __init__(self, match):
super().__init__(match) super().__init__(match)
self.params.update(text.parse_query(match.group(2))) self.params = text.parse_query(match.group(2))
def metadata(self): def metadata(self):
self.page_start = self.page_start or 1
scale = self.params.get("scale", "day") scale = self.params.get("scale", "day")
date = self.params.get("date") or datetime.date.today().isoformat() date = self.params.get("date") or datetime.date.today().isoformat()
@@ -209,5 +216,7 @@ class DanbooruPopularExtractor(DanbooruExtractor):
return {"date": date, "scale": scale} return {"date": date, "scale": scale}
def posts(self): def posts(self):
url = self.root + "/explore/posts/popular.json" if self.page_start is None:
return self._pagination(url, True) self.page_start = 1
return self._pagination(
"/explore/posts/popular.json", self.params, True)