Refactored to make fetching and downloading posts seperate

This commit is contained in:
Anon 2022-08-29 15:15:09 -07:00
parent 8fb4da20ad
commit bfaaa50117

View File

@ -24,24 +24,33 @@ class downloader:
username = None username = None
password = None password = None
tmp = None tmp = None
banned = None unallowed_extensions = (".zip",)
def __init__(self, banned = tuple(), username=None, password=None, tmp="/tmp"): def __init__(self, username=None, password=None, tmp="/tmp"):
self.username = username self.username = username
self.password = password self.password = password
self.tmp = tmp self.tmp = tmp
self.banned = banned
def is_banned(self, tag_list): def download_post(self, post):
for tag in self.banned: file_url = post["file_url"]
if tag in tag_list: full_path = post["full_path"]
return True
return False remote_image = requests.get(file_url)
if remote_image.status_code != 200:
print("Remote image request returned:", remote_image.status_code)
return None
for d in full_path:
with open(d, "wb") as f:
f.write(remote_image.content)
return post
def download(self, profile): def fetch_post(self, profile):
# Search ratings: s=safe, e=explicit # 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=1"
@ -50,7 +59,6 @@ class downloader:
search_url = "{}&tags={}".format(search_url, search_tags) search_url = "{}&tags={}".format(search_url, search_tags)
search_request = None search_request = None
while True:
if self.username and self.password: if self.username and self.password:
search_request = requests.get(search_url, search_request = requests.get(search_url,
auth=(self.username, self.password) auth=(self.username, self.password)
@ -60,40 +68,25 @@ class downloader:
if search_request.status_code != 200: if search_request.status_code != 200:
print("Search request returned:", search_request.status_code) print("Search request returned:", search_request.status_code)
continue return None
elif "large_file_url" not in search_request.json()[0]: response = search_request.json()[0]
continue tag_response = []
elif "tag_string" not in search_request.json()[0]: if "file_url" not in response:
continue print("file_url is not in response")
elif "tag_string_general" not in search_request.json()[0]:
continue
elif self.is_banned(search_request.json()[0]["tag_string"]):
print("Banned Tag1:", search_request.json()[0]["tag_string"])
continue
elif self.is_banned(search_request.json()[0]["tag_string_general"]):
print("Banned Tag2",search_request.json()[0]["tag_string_general"])
continue
break
import pprint
pprint.pprint(search_request.json()[0]["tag_string"])
large_file_url = search_request.json()[0]["file_url"]
explicit = search_request.json()[0]["rating"]
explicit = get_most_sever_rating(explicit)
remote_image = requests.get(large_file_url)
if remote_image.status_code != 200:
print("Remote image request returned:", remote_image.status_code)
return None return None
basename = large_file_url.rsplit("/", 1)[1] # Aggregate Tags
full_path = os.path.join(self.tmp, basename) for tag_type in "tag_string", "tag_string_general":
if tag_type in response:
tag_response.append(response[tag_type].strip())
with open(full_path, "wb") as f: nsfw = search_request.json()[0]["rating"]
f.write(remote_image.content) nsfw = get_most_sever_rating(nsfw)
file_url = response["file_url"]
basename = file_url.rsplit("/", 1)[1]
full_path = os.path.join(self.tmp, basename)
r = { r = {
# Add profile to dictioanry # Add profile to dictioanry
@ -101,13 +94,14 @@ class downloader:
"backend": profile["backend"], "backend": profile["backend"],
"tags": profile["tags"], "tags": profile["tags"],
"message": profile["message"], "message": profile["message"],
"message_nsfw": profile["message_nsfw"],
# Query results # Query results
"search_url": search_url, "search_url": search_url,
"large_file_url": large_file_url, "file_url": file_url,
"full_path": [full_path], "full_path": [full_path],
"explicit": explicit, "tag_response": " ".join(tag_response),
"nsfw": nsfw
} }
return r return r