Fetches multiple posts at once, improves performance
This commit is contained in:
parent
79c099cdf1
commit
8b31de29f6
@ -4,6 +4,7 @@
|
||||
# Update to enable multiple file downloads
|
||||
# Update to include safe/nsfw filter
|
||||
|
||||
import magic
|
||||
import requests
|
||||
import os
|
||||
|
||||
@ -12,15 +13,46 @@ def random_tag(*tags):
|
||||
return len(tags) == 1 and tags[0].lower() == "random"
|
||||
|
||||
|
||||
def get_most_sever_rating(rating):
|
||||
return rating in ("q", "e")
|
||||
def collect_tags(post):
|
||||
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:
|
||||
username = None
|
||||
password = None
|
||||
tmp = None
|
||||
unallowed_extensions = (".zip",)
|
||||
limit=50
|
||||
|
||||
def __init__(self, username=None, password=None, tmp="/tmp"):
|
||||
self.username = username
|
||||
@ -48,7 +80,7 @@ class downloader:
|
||||
# Search ratings: s=safe, e=nsfw
|
||||
# base_url = "https://danbooru.donmai.us/posts.json?random=true&tags={}&rating=e&limit=1"
|
||||
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):
|
||||
search_tags = "+".join(tags)
|
||||
search_url = "{}&tags={}".format(search_url, search_tags)
|
||||
@ -62,34 +94,32 @@ class downloader:
|
||||
search_request = requests.get(search_url)
|
||||
|
||||
if search_request.status_code != 200:
|
||||
print(search_url)
|
||||
print("Search request returned:", search_request.status_code)
|
||||
return None
|
||||
response = search_request.json()[0]
|
||||
tag_response = []
|
||||
if "file_url" not in response:
|
||||
print("file_url is not in response")
|
||||
return None
|
||||
|
||||
# Aggregate Tags
|
||||
for tag_type in "tag_string", "tag_string_general":
|
||||
if tag_type in response:
|
||||
tag_response.append(response[tag_type].strip())
|
||||
selected = select_from_response(search_request.json(), profile)
|
||||
|
||||
nsfw = response["rating"]
|
||||
nsfw = get_most_sever_rating(nsfw)
|
||||
if selected is None:
|
||||
print("Could not select image based on criteria")
|
||||
return None
|
||||
|
||||
file_url = response["file_url"]
|
||||
tag_response = collect_tags(selected)
|
||||
nsfw = get_nsfw(selected)
|
||||
file_url = selected["file_url"]
|
||||
|
||||
basename = file_url.rsplit("/", 1)[1]
|
||||
full_path = os.path.join(self.tmp, basename)
|
||||
|
||||
r = {
|
||||
# Add profile to dictioanry
|
||||
"name": profile["name"],
|
||||
"backend": profile["backend"],
|
||||
"tags": profile["tags"],
|
||||
"message": profile["message"],
|
||||
"message_nsfw": profile["message_nsfw"],
|
||||
"name": profile["name"],
|
||||
"backend": profile["backend"],
|
||||
"tags": profile["tags"],
|
||||
"message": profile["message"],
|
||||
"message_nsfw": profile["message_nsfw"],
|
||||
"force_nsfw": profile["force_nsfw"],
|
||||
"banned_tags": profile["banned_tags"],
|
||||
|
||||
# Query results
|
||||
"search_url": search_url,
|
||||
|
Reference in New Issue
Block a user