feat: save generated pictures to a cache per seed

pfp_generator/cli.py

@@ -8,6 +8,9 @@
 from pfp_generator.generate import ColorMatrix
 
 
+MAX_PATTERN_SIZE: int = 100
+
+
 def clean_color(color: str, seed: str) -> colors.RGB:
     """Sanitize the color string to a valid RGB color.
 
@@ -96,8 +99,7 @@ def main() -> None:
         type=str,
         nargs="?",
         default=None,
-        help="Text to generate the profile picture from. If left blank, generate a "
-        "random profile picture.",
+        help="Text to generate from. Random text is used if not provided.",
     )
 
     parser.add_argument(
@@ -147,6 +149,9 @@ def main() -> None:
     )
     args = parser.parse_args()
 
+    if args.size >= MAX_PATTERN_SIZE:
+        return print(f"Pattern size must be less than {MAX_PATTERN_SIZE}!")
+
     default_seed: int | str = clean_seed(str(time.time()).replace(".", ""))
     seed: int | str = clean_seed(args.text) if args.text else default_seed
     color_seed: int | str = clean_seed(str(seed)[::-1])

pfp_generator/display.py

@@ -1,12 +1,14 @@
 """Display a profile picture pattern with a given color matrix."""
 
-from tkinter.filedialog import asksaveasfilename
+from pathlib import Path
 
 from PIL import Image
 
 from pfp_generator.generate import ColorMatrix, batch_generate_pfp, generate_pfp
 
 PER_ROW: int = 5
+SAVE_PATH: Path = Path(Path("~").expanduser(), ".cache", "pfp-generator")
+CACHE_LIMIT: int = 5  # MB
 
 
 def display_pfp(
@@ -23,7 +25,7 @@ def display_pfp(
         save (bool): A flag to save the profile picture. Defaults to False.
     """
     pattern = [generate_pfp(matrices[0])]
-    name = "pfps" if len(matrices) > 1 else matrices[0].seed
+    name = str(matrices[0].seed)
 
     pattern.extend(batch_generate_pfp(matrices[1:]))
 
@@ -35,31 +37,71 @@ def display_pfp(
 
     collage = Image.new("RGB", (collage_width, collage_height))
 
+    images = []
+
     for i, p in enumerate(pattern):
         image = (
             Image.fromarray(p.astype("uint8"))
             .convert("RGB")
             .resize((size, size), Image.NEAREST)
         )
+        images.append(image)
         row, col = divmod(i, PER_ROW)
         collage.paste(image, (col * size, row * size))
 
     collage.show()
+
+    if cache_exceeded():
+        return print(
+            f"Cache limit reached! If you want to save more images, please clear the "
+            f"cache @ {SAVE_PATH}"
+        )
+
     if save:
-        save_file(collage, name)
+        save_file(images, name)
+
+
+def cache_exceeded() -> bool:
+    """Check if the cache exceeds the cache limit.
+
+    Returns:
+        bool: True if the cache exceeds the cache limit, False otherwise.
+    """
+    return get_dir_size(SAVE_PATH) >= CACHE_LIMIT
+
+
+def get_dir_size(directory: Path) -> int:
+    """Get the size of a directory and its subdirectories in MB.
+
+    Args:
+        directory (Path): The directory to get the size of.
+
+    Returns:
+        int: The size of the directory in bytes.
+    """
+    size_bytes = sum(f.stat().st_size for f in directory.glob("**/*") if f.is_file())
+    return size_bytes / 1024 / 1024
 
 
-def save_file(image: Image.Image, name: str) -> None:
+def save_file(images: list[Image.Image], name: str) -> None:
     """Save an image file.
 
     Args:
-        image (Image.Image): The image to save.
+        images (list[Image.Image]): The images to save.
         name (str): The name of the image.
     """
-    filename = asksaveasfilename(
-        filetypes=[("PNG files", "*.png"), ("JPEG files", "*.jpg")],
-        defaultextension=".png",
-        initialfile=name,
-    )
-    if filename:
-        image.save(filename)
+    save_dir = Path(SAVE_PATH, name)
+    if not Path.exists(save_dir):
+        Path.mkdir(save_dir, parents=True)
+
+    for i, image in enumerate(images, start=1):
+        if cache_exceeded():
+            return print(
+                f"Cache limit reached! If you want to save more images, please clear "
+                f"the cache @ {SAVE_PATH}"
+            )
+        filename = Path(save_dir, f"{name}_{i}.png")
+        if not Path.exists(filename):
+            image.save(filename)
+
+    print(f"Profile pictures saved to {Path(SAVE_PATH, name)}")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pfp-generator"
-version = "0.3.4"
+version = "0.4.0"
 description = "Generate profile pictures."
 authors = ["trumully <truman.mulholland@gmail.com>"]
 readme = "README.md"