Fetches multiple posts at once, improves performance

src/danbooru_backend.py

@@ -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):
+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:
 	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,23 +94,19 @@ 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 = []
+		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
 
-		# 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]
 		full_path = os.path.join(self.tmp, basename)
@@ -90,6 +118,8 @@ class downloader:
 			"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,