StaticGalleryBuilder/builder.py

#!/usr/bin/env python3
import os
import re
import shutil
import fnmatch
import urllib.parse
from multiprocessing import Pool
from pathlib import Path
from typing import Dict, List, Tuple

from tqdm.auto import tqdm
from PIL import Image, ImageOps

from modules.argumentparser import parse_arguments, Args
from modules.svg_handling import icons, webmanifest, extract_colorscheme
from modules.generate_html import list_folder, EXCLUDES

# fmt: off
# Constants
STATIC_FILES_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), "files")
SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))
VERSION = "2.2.3"
RAW_EXTENSIONS = [
    ".3fr", ".ari", ".arw", ".bay", ".braw", ".crw", ".cr2", ".cr3", ".cap", ".data", ".dcs", ".dcr",
    ".dng", ".drf", ".eip", ".erf", ".fff", ".gpr", ".iiq", ".k25", ".kdc", ".mdc", ".mef", ".mos",
    ".mrw", ".nef", ".nrw", ".obm", ".orf", ".pef", ".ptx", ".pxn", ".r3d", ".raf", ".raw", ".rwl",
    ".rw2", ".rwz", ".sr2", ".srf", ".srw", ".tif", ".tiff", ".x3f"
]
IMG_EXTENSIONS = [".jpg", ".jpeg", ".png"]
NOT_LIST = ["*/Galleries/*", "Archives"]
# fmt: on

pbardict: Dict[str, tqdm] = {}


def init_globals(_args: Args, raw: List[str]) -> Tuple[Args, List[str]]:
    """
    Initialize global variables and set default values for arguments.

    Parameters:
    -----------
    _args : Args
        Parsed command-line arguments.
    raw : List[str]
        List of raw file extensions.

    Returns:
    --------
    Tuple[Args, List[str]]
        Updated arguments and raw file extensions.
    """
    if not _args.file_extensions:
        _args.file_extensions = IMG_EXTENSIONS
    if not _args.exclude_folders:
        _args.exclude_folders = NOT_LIST
    _args.root_directory = _args.root_directory.rstrip("/") + "/"
    _args.web_root_url = _args.web_root_url.rstrip("/") + "/"

    raw = [ext.lower() for ext in raw] + [ext.upper() for ext in raw]

    return _args, raw


def copy_static_files(_args: Args) -> None:
    """
    Copy static files to the root directory.

    Parameters:
    -----------
    _args : Args
        Parsed command-line arguments.
    """
    static_dir = os.path.join(_args.root_directory, ".static")
    if os.path.exists(static_dir):
        print("Removing existing .static folder...")
        shutil.rmtree(static_dir)

    print("Copying static files...")
    shutil.copytree(STATIC_FILES_DIR, static_dir, dirs_exist_ok=True)
    with open(_args.theme_path, "r", encoding="utf-8") as f:
        theme = f.read()
    split = theme.split(".foldericon {")
    split2 = split[1].split("}", maxsplit=1)
    themehead = split[0]
    themetail = split2[1]
    foldericon = split2[0].strip()
    foldericon = re.sub(r"/\*.*\*/", "", foldericon)
    for match in re.finditer(r"content: (.*);", foldericon):
        foldericon = match[1]
        foldericon = foldericon.replace('"', "")
        break
    if "url" in foldericon:
        shutil.copyfile(_args.theme_path, os.path.join(static_dir, "theme.css"))
        return
    with open(os.path.join(SCRIPT_DIR, foldericon), "r", encoding="utf-8") as f:
        svg = f.read()
    if "svg.j2" in foldericon:
        colorscheme = extract_colorscheme(_args.theme_path)
        svg = svg.replace("{{ color1 }}", colorscheme["color1"])
        svg = svg.replace("{{ color2 }}", colorscheme["color2"])
        svg = svg.replace("{{ color3 }}", colorscheme["color3"])
        svg = svg.replace("{{ color4 }}", colorscheme["color4"])
    svg = urllib.parse.quote(svg)
    with open(os.path.join(static_dir, "theme.css"), "x", encoding="utf-8") as f:
        f.write(themehead + '\n.foldericon {\n  content: url("data:image/svg+xml,' + svg + '");\n}\n' + themetail)


def generate_thumbnail(arguments: Tuple[str, str, str]) -> None:
    """
    Generate a thumbnail for a given image.

    Parameters:
    -----------
    arguments : Tuple[str, str, str, bool]
        A tuple containing the folder, item, root directory, and regenerate thumbnails flag.
    """
    folder, item, root_directory = arguments
    path = os.path.join(root_directory, ".thumbnails", folder.removeprefix(root_directory), item) + ".jpg"
    oldpath = os.path.join(root_directory, ".thumbnails", folder.removeprefix(root_directory), os.path.splitext(item)[0]) + ".jpg"
    if os.path.exists(oldpath):
        try:
            shutil.move(oldpath, path)
        except FileNotFoundError:
            pass
    if not os.path.exists(path):
        try:
            with Image.open(os.path.join(folder, item)) as imgfile:
                imgrgb = imgfile.convert("RGB")
                img = ImageOps.exif_transpose(imgrgb)
                img.thumbnail((512, 512))
                img.save(path, "JPEG", quality=75, optimize=True, mode="RGB")
        except OSError:
            print(f"Failed to generate thumbnail for {os.path.join(folder, item)}")


def get_total_folders(folder: str, _args: Args, _total: int = 0) -> int:
    """
    Recursively count the total number of folders to be processed.

    Parameters:
    -----------
    folder : str
        The current folder being processed.
    _args : Args
        Parsed command-line arguments.
    _total : int, optional
        The running total of folders, default is 0.

    Returns:
    --------
    int
        The total number of folders.
    """
    _total += 1
    pbardict["traversingbar"].desc = f"Traversing filesystem - {folder}"
    pbardict["traversingbar"].update(1)

    items = sorted(os.listdir(folder))
    for item in items:
        if item not in EXCLUDES and os.path.isdir(os.path.join(folder, item)):
            if item not in _args.exclude_folders and not any(
                fnmatch.fnmatchcase(os.path.join(folder, item), exclude) for exclude in _args.exclude_folders
            ):
                _total = get_total_folders(os.path.join(folder, item), _args, _total)
    return _total


def main() -> None:
    """
    Main function to process images and generate a static image hosting website.
    """
    thumbnails: List[Tuple[str, str, str, bool]] = []

    args = parse_arguments(VERSION)
    args, raw = init_globals(args, RAW_EXTENSIONS)

    lock_file = os.path.join(args.root_directory, ".lock")
    if os.path.exists(lock_file):
        print("Another instance of this program is running.")
        exit()

    try:
        Path(lock_file).touch()
        os.makedirs(os.path.join(args.root_directory, ".thumbnails"), exist_ok=True)

        copy_static_files(args)
        icons(args)

        if args.generate_webmanifest:
            print("Generating webmanifest...")
            webmanifest(args)

        if args.non_interactive_mode:
            print("Generating HTML files...")
            thumbnails = list_folder(0, args.root_directory, args.site_title, args, raw, VERSION)
            with Pool(os.cpu_count()) as pool:
                print("Generating thumbnails...")
                pool.map(generate_thumbnail, thumbnails)
        else:
            pbardict["traversingbar"] = tqdm(desc="Traversing filesystem", unit="folders", ascii=True, dynamic_ncols=True)
            total = get_total_folders(args.root_directory, args)
            pbardict["traversingbar"].desc = "Traversing filesystem"
            pbardict["traversingbar"].update(0)
            pbardict["traversingbar"].close()

            thumbnails = list_folder(total, args.root_directory, args.site_title, args, raw, VERSION)

            with Pool(os.cpu_count()) as pool:
                for _ in tqdm(
                    pool.imap_unordered(generate_thumbnail, thumbnails),
                    total=len(thumbnails),
                    desc="Generating thumbnails",
                    unit="files",
                    ascii=True,
                    dynamic_ncols=True,
                ):
                    pass
    finally:
        os.remove(lock_file)
    return


if __name__ == "__main__":
    main()