Fetches multiple posts at once, improves performance

This commit is contained in:
Anon 2022-08-30 15:17:27 -07:00
parent 79c099cdf1
commit 8b31de29f6

View File

@ -4,6 +4,7 @@
# Update to enable multiple file downloads # Update to enable multiple file downloads
# Update to include safe/nsfw filter # Update to include safe/nsfw filter
import magic
import requests import requests
import os import os
@ -12,15 +13,46 @@ def random_tag(*tags):
return len(tags) == 1 and tags[0].lower() == "random" return len(tags) == 1 and tags[0].lower() == "random"
def get_most_sever_rating(rating): def collect_tags(post):
return rating in ("q", "e") tag_response = []
for tag_type in "tag_string", "tag_string_general":
if tag_type in post:
tag_response.append(post[tag_type].strip())
return " ".join(tag_response)
def is_banned(post, profile):
tag_response = collect_tags(post)
tag_banned = profile["banned_tags"]
for tag in tag_banned:
if tag in tag_response:
return tag
return None
def get_nsfw(post):
return post["rating"] in ("q", "e")
def select_from_response(response, profile):
for post in response:
if is_banned(post, profile):
continue
elif "file_url" not in post:
continue
elif profile["force_nsfw"] is not None and profile["force_nsfw"] != get_nsfw(post):
continue
return post
return None
class downloader: class downloader:
username = None username = None
password = None password = None
tmp = None tmp = None
unallowed_extensions = (".zip",) limit=50
def __init__(self, username=None, password=None, tmp="/tmp"): def __init__(self, username=None, password=None, tmp="/tmp"):
self.username = username self.username = username
@ -48,7 +80,7 @@ class downloader:
# Search ratings: s=safe, e=nsfw # Search ratings: s=safe, e=nsfw
# base_url = "https://danbooru.donmai.us/posts.json?random=true&tags={}&rating=e&limit=1" # base_url = "https://danbooru.donmai.us/posts.json?random=true&tags={}&rating=e&limit=1"
tags = profile["tags"] tags = profile["tags"]
search_url = "https://danbooru.donmai.us/posts.json?random=true&limit=1" search_url = "https://danbooru.donmai.us/posts.json?random=true&limit={}".format(self.limit)
if tags and not random_tag(*tags): if tags and not random_tag(*tags):
search_tags = "+".join(tags) search_tags = "+".join(tags)
search_url = "{}&tags={}".format(search_url, search_tags) search_url = "{}&tags={}".format(search_url, search_tags)
@ -62,23 +94,19 @@ class downloader:
search_request = requests.get(search_url) search_request = requests.get(search_url)
if search_request.status_code != 200: if search_request.status_code != 200:
print(search_url)
print("Search request returned:", search_request.status_code) print("Search request returned:", search_request.status_code)
return None return None
response = search_request.json()[0]
tag_response = [] selected = select_from_response(search_request.json(), profile)
if "file_url" not in response:
print("file_url is not in response") if selected is None:
print("Could not select image based on criteria")
return None return None
# Aggregate Tags tag_response = collect_tags(selected)
for tag_type in "tag_string", "tag_string_general": nsfw = get_nsfw(selected)
if tag_type in response: file_url = selected["file_url"]
tag_response.append(response[tag_type].strip())
nsfw = response["rating"]
nsfw = get_most_sever_rating(nsfw)
file_url = response["file_url"]
basename = file_url.rsplit("/", 1)[1] basename = file_url.rsplit("/", 1)[1]
full_path = os.path.join(self.tmp, basename) full_path = os.path.join(self.tmp, basename)
@ -90,6 +118,8 @@ class downloader:
"tags": profile["tags"], "tags": profile["tags"],
"message": profile["message"], "message": profile["message"],
"message_nsfw": profile["message_nsfw"], "message_nsfw": profile["message_nsfw"],
"force_nsfw": profile["force_nsfw"],
"banned_tags": profile["banned_tags"],
# Query results # Query results
"search_url": search_url, "search_url": search_url,