Compare commits

...

2 Commits

2 changed files with 20 additions and 11 deletions

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")

View File

@ -144,12 +144,10 @@ class YandereBot:
try: try:
backend_s = picked_profile["backend"] backend_s = picked_profile["backend"]
backend = importlib.import_module(backend_s) backend = importlib.import_module(backend_s)
username = self.settings_credentials[backend_s]["username"] backend_credentials = self.settings_credentials[backend_s]
password = self.settings_credentials[backend_s]["password"]
img = None img = None
downloader = backend.downloader( downloader = backend.downloader(
username, password, backend_credentials, self.settings_behavior
self.settings_behavior["tmp_dir"]
) )
img = downloader.fetch_post(picked_profile) img = downloader.fetch_post(picked_profile)