From 7d398cc82411b24860c6e507565da933bc9d00fd Mon Sep 17 00:00:00 2001
From: Kat Inskip <kat@inskip.me>
Date: Fri, 28 Feb 2025 00:04:32 -0800
Subject: [PATCH] feat: add e621

---
 konawall/gui.py              |  4 ++-
 konawall/sources/e621.py     | 66 ++++++++++++++++++++++++++++++++++++
 konawall/sources/konachan.py |  4 +--
 3 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 konawall/sources/e621.py

diff --git a/konawall/gui.py b/konawall/gui.py
index baf8c07..9f155fa 100755
--- a/konawall/gui.py
+++ b/konawall/gui.py
@@ -24,6 +24,7 @@ class Konawall(wx.adv.TaskBarIcon):
             self.hidden_frame.Hide()
 
         self.log_path = log_path
+        self.config = {}
         self.wallpaper_rotation_counter = 0 
         self.file_logger = file_logger
         self.version = version
@@ -136,6 +137,7 @@ class Konawall(wx.adv.TaskBarIcon):
             # If the config file exists, load it as a dictionary into the config variable.
             with open(self.config_path, "rb") as f:
                 config = tomllib.load(f)
+                self.config = config
             # for every key-value pair in the config variable , set the corresponding attribute of our class to it
             for k, v in config.items():
                 kv_print(f"Loaded {k}", v)
@@ -341,7 +343,7 @@ class Konawall(wx.adv.TaskBarIcon):
     def rotate_wallpapers(self, event):
         displays = screeninfo.get_monitors()
         count = len(displays)
-        files, self.current = source_handlers[self.source](count, self.tags)
+        files, self.current = source_handlers[self.source](count, self.tags, self.config)
         set_environment_wallpapers(self.environment, files, displays)
 
     # For macOS
diff --git a/konawall/sources/e621.py b/konawall/sources/e621.py
new file mode 100644
index 0000000..eae8230
--- /dev/null
+++ b/konawall/sources/e621.py
@@ -0,0 +1,66 @@
+import requests
+import logging
+from konawall.custom_print import kv_print
+from konawall.custom_errors import RequestFailed
+from konawall.module_loader import add_source
+from konawall.downloader import download_files
+
+"""
+Turn a list of tags and a count into a list of URLs to download from
+
+:param count: The number of images to provide download URLs for
+:param user_tags: A list of tags to search for
+:returns: A list of URLs to download from
+"""
+def request_posts(count: int, tags: list, config) -> list:
+    api_key = config["api_key"]
+    logging.debug(f"request_posts() called with count={count}, tags=[{', '.join(tags)}]")
+    # Make sure we get a different result every time by using "order:random" as a tag
+    if "order:random" not in tags:
+        tags.append("order:random")
+    # Tags are separated by a plus sign for this API
+    tag_string: str = "+".join(tags)
+    # Request URL for getting posts from the API
+    url: str = f"https://e621.net/posts.json?limit={str(count)}&tags={tag_string}"
+    logging.debug(f"Request URL: {url}")
+    response = requests.get(url, headers={"User-Agent": "konachan-py/alpha (by katsmew on e621)"})
+    # Check if the request was successful
+    logging.debug("Status code: " + str(response.status_code))
+    # List of URLs to download
+    posts: list = []
+    if response.status_code == 200:
+        # Get the JSON data from the response
+        json = response.json()
+        for post in json["posts"]:
+            # Give the user data about the post retrieved
+            kv_print("Post ID", post["id"])
+            kv_print("Author", post["uploader_id"])
+            kv_print("Rating", post["rating"])
+            kv_print("Resolution", f"{post['file']['width']}x{post['file']['height']}")
+            kv_print("Tags", post["tags"])
+            kv_print("URL", post["file"]["url"])
+            # Append the URL to the list
+            posts.append(post)
+    else:
+        # Raise an exception if the request failed
+        RequestFailed(response.status_code)
+    return posts
+
+"""
+Download a number of images from Konachan given a list of tags and a count
+
+:param count: The number of images to download
+:param tags: A list of tags to search for
+"""
+@add_source("e621")
+def handle(count: int, tags: list, config) -> list:
+    logging.debug(f"handle_e621() called with count={count}, tags=[{', '.join(tags)}]")
+    # Get a list of URLs to download
+    posts: list = request_posts(count, tags, config)
+    urls: list = []
+    # Download the images
+    for post in posts:
+        urls.append(post["file"]["url"])
+    files = download_files(urls)
+    # Return the downloaded files
+    return files, posts
diff --git a/konawall/sources/konachan.py b/konawall/sources/konachan.py
index 4f5a357..60035db 100644
--- a/konawall/sources/konachan.py
+++ b/konawall/sources/konachan.py
@@ -12,7 +12,7 @@ Turn a list of tags and a count into a list of URLs to download from
 :param user_tags: A list of tags to search for
 :returns: A list of URLs to download from
 """
-def request_posts(count: int, tags: list) -> list:
+def request_posts(count: int, tags: list, config={}) -> list:
     logging.debug(f"request_posts() called with count={count}, tags=[{', '.join(tags)}]")
     # Make sure we get a different result every time by using "order:random" as a tag
     if "order:random" not in tags:
@@ -62,4 +62,4 @@ def handle(count: int, tags: list) -> list:
         urls.append(post["file_url"])
     files = download_files(urls)
     # Return the downloaded files
-    return files, posts
\ No newline at end of file
+    return files, posts