Anon/DanbooruBot
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added function to validate file size before returning post

Browse Source
This commit is contained in:
Anon 2022-09-01 18:58:09 -07:00
parent 4b466f206c
commit ed44b50414
1 changed files with 18 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
src/danbooru_backend.py
View File

@ -50,13 +50,22 @@ def get_nsfw(post):
return post["rating"] in ("q", "e") return post["rating"] in ("q", "e")
def select_from_response(response, profile): def select_from_response(response, profile, max_size=None):
for post in response: for post in response:
if is_banned(post, profile): if is_banned(post, profile):
continue continue
elif "file_url" not in post: elif "file_url" not in post:
continue continue
elif profile["force_nsfw"] is not None and profile["force_nsfw"] != get_nsfw(post): # Select only nsfw
elif ( profile["force_nsfw"] is not None and
profile["force_nsfw"] != get_nsfw(post)
):
continue
# Make sure serverside size is not larger than max_size
elif ( max_size != None and
"file_size" in post and
post["file_size"] > max_size
):
continue continue
return post return post
return None return None
@ -65,14 +74,16 @@ def select_from_response(response, profile):
class downloader: class downloader:
username = None username = None
password = None password = None
max_size = None
tmp = None tmp = None
# Limit for posts.json is 200 # Limit for posts.json is 200
limit=100 limit=100
def __init__(self, username=None, password=None, tmp="/tmp"): def __init__(self, backend_credentials, settings_behavior):
self.username = username self.username = backend_credentials["username"]
self.password = password self.password = backend_credentials["password"]
self.tmp = tmp self.max_size = settings_behavior["max_size"]
self.tmp = settings_behavior["tmp_dir"]
def download_post(self, post): def download_post(self, post):
@ -113,7 +124,7 @@ class downloader:
print("Search request returned:", search_request.status_code) print("Search request returned:", search_request.status_code)
return None return None
selected = select_from_response(search_request.json(), profile) selected = select_from_response(search_request.json(), profile, self.max_size)
if selected is None: if selected is None:
print("Could not select image based on criteria") print("Could not select image based on criteria")