walthemer.py

#!/usr/bin/env python3

import os
import subprocess
import argparse

from wallhavenapi import WallhavenAPI, Purity, Sorting, Category

def enum_args(parser, name: str, enum_cls, help: str):
    parser.add_argument(
        f"--{name}",
        type=lambda v: enum_cls(v),
        choices=list(enum_cls),
        nargs="+",
        help=f"{help} (choices: {[e.value for e in enum_cls]})"
    )

def resolution_type(value: str):
    try:
        w, h = value.lower().split("x")
        return (int(w), int(h))
    except ValueError:
        raise argparse.ArgumentTypeError(f"Resolution must be in WxH format, e.g. 1920x1080 (got '{value}')")

def parse_args():
    parser = argparse.ArgumentParser(
        description="Walthemer - Download a wallpaper from wallhaven and create a colorscheme using pywal16."
    )
    
    parser.add_argument(
        "--api-key",
        type=str,
        help="Path of the file storing your API key"
    )
    
    parser.add_argument(
        "--query",
        type=str,
        help="Keyword used to search Wallhaven"
    )
    
    parser.add_argument(
        "--output",
        type=str,
        default="wallpaper.jpg",
        help="File path to save the wallpaper to (Default: ~/wallpaper.jpg)"
    )
    
    parser.add_argument(
        "--resolution",
        type=resolution_type,
        help="Wallpaper resolution in the format WxH (e.g. 1920x1080)."
    )
    
    enum_args(
        parser,
        "purities",
        Purity,
        "Preferred purities to apply in search; 'nsfw' and 'sketchy' only usable if API key is provided"
    )
    
    enum_args(
        parser,
        "sorting",
        Sorting,
        "Preferred sorting to apply in search"
    )
    
    enum_args(
        parser,
        "categories",
        Category,
        "Preferred categories to search in"
    )
    
    args = parser.parse_args()
    
    if not args.api_key:
        if args.purities and any(p != Purity.sfw for p in args.purities):
            parser.error("Using 'sketchy' or 'nsfw' purities requires --api-key")
        
        if not args.purities:
            args.purities = [Purity.sfw]
    
    for name, enum_cls in {"purities": Purity, "sorting": Sorting, "categories": Category}.items():
        values = getattr(args, name)
        if values and len(values) > len(enum_cls):
            parser.error(f"--{name} accepts at most {len(enum_cls)} values.")
            
    return args

def download_wallpaper(
    api_key=None,
    query=None,
    save_path=None,
    resolution=None,
    purities=None,
    sorting=None,
    categories=None
):
    api = WallhavenAPI(api_key=api_key)

    if os.path.isabs(save_path) or save_path.startswith("~"):
        save_path = os.path.expanduser(save_path)
    else:
        save_path = os.path.join(os.path.expanduser("~"), save_path)
    
    result = api.search(
        q=query,
        categories=categories,
        purities=purities,
        sorting=sorting,
        atleast=resolution
    )
    
    wallpapers = result.get("data", [])
    
    if not wallpapers:
        print("No wallpapers found.")
        return
        
    wallpaper_id = wallpapers[0]["id"]
    
    print(f"Downloading wallpaper with ID: {wallpaper_id}")
    path = api.download_wallpaper(wallpaper_id=wallpaper_id, file_path=save_path)
    if path:
        print(f"Wallpaper saved to {path}")
        return path

def generate_colorscheme(wallpaper_path):
    try:
        subprocess.run(["wal", "-i", wallpaper_path], check=True)
        print("Colorscheme generated using pywal16!")
    except FileNotFoundError:
        print("Error: 'wal16' (pywal16) is not installed or not in PATH.")
    except subprocess.CalledProcessError:
        print("Error: Failed to generate colorscheme with pywal16.")


def set_wallpaper(wallpaper_path):
    try:
        subprocess.run(["swww", "img", wallpaper_path], check=True)
    except FileNotFoundError:
        print("Error: 'swww' is not installed.")
    except subprocess.CalledProcessError:
        print("Error: Failed to set wallpaper using swww. Is the daemon running?")

def main():
    args = parse_args()
    
    # Replace with your Wallhaven API key, or leave 'None' for public access (SFW only)
    api_key = args.api_key or None
    query = args.query or None
    output = args.output
    resolution = args.resolution or None
    purities = args.purities
    sorting = args.sorting or None
    categories = args.categories or None

    wallpaper_path = download_wallpaper(
        api_key,
        query,
        output,
        resolution,
        purities,
        sorting,
        categories
    )
    
    if wallpaper_path:
        generate_colorscheme(wallpaper_path)
        set_wallpaper(wallpaper_path)

if __name__ == "__main__":
    main()