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

View File

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