diff --git a/src/danbooru_backend.py b/src/danbooru_backend.py index 06daff9..236de75 100644 --- a/src/danbooru_backend.py +++ b/src/danbooru_backend.py @@ -50,13 +50,22 @@ def get_nsfw(post): 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: 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): + # 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 return post return None @@ -65,14 +74,16 @@ def select_from_response(response, profile): class downloader: username = None password = None + max_size = None tmp = None # Limit for posts.json is 200 limit=100 - def __init__(self, username=None, password=None, tmp="/tmp"): - self.username = username - self.password = password - self.tmp = tmp + def __init__(self, backend_credentials, settings_behavior): + self.username = backend_credentials["username"] + self.password = backend_credentials["password"] + self.max_size = settings_behavior["max_size"] + self.tmp = settings_behavior["tmp_dir"] def download_post(self, post): @@ -113,7 +124,7 @@ class downloader: print("Search request returned:", search_request.status_code) 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: print("Could not select image based on criteria")