separated into modules and documented everything

modules/generate_html.py

@@ -0,0 +1,313 @@
+import os
+import urllib.parse
+import fnmatch
+import json
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+from tqdm.auto import tqdm
+from PIL import Image, ExifTags
+from jinja2 import Environment, FileSystemLoader
+
+import modules.cclicense as cclicense
+from modules.argumentparser import Args
+
+# Constants for file paths and exclusions
+FAVICON_PATH = ".static/favicon.ico"
+GLOBAL_CSS_PATH = ".static/global.css"
+EXCLUDES = [".lock", "index.html", "manifest.json", ".sizelist.json", ".thumbnails", ".static"]
+
+# Set the maximum image pixels to prevent decompression bomb DOS attacks
+Image.MAX_IMAGE_PIXELS = 933120000
+
+# Initialize Jinja2 environment for template rendering
+env = Environment(loader=FileSystemLoader(os.path.join(os.path.abspath(os.path.dirname(__file__)), "..", "templates")))
+thumbnails: List[Tuple[str, str]] = []
+info: Dict[str, str] = {}
+pbardict: Dict[str, tqdm] = {}
+
+
+def initialize_sizelist(folder: str) -> Dict[str, Dict[str, int]]:
+    """
+    Initializes the size list JSON file if it doesn't exist.
+
+    Args:
+        folder (str): The folder in which the size list file is located.
+
+    Returns:
+        Dict[str, Dict[str, int]]: The size list dictionary.
+    """
+    sizelist = {}
+    sizelist_path = os.path.join(folder, ".sizelist.json")
+    if not os.path.exists(sizelist_path):
+        with open(sizelist_path, "x", encoding="utf-8") as sizelistfile:
+            sizelistfile.write("{}")
+    with open(sizelist_path, "r+", encoding="utf-8") as sizelistfile:
+        try:
+            sizelist = json.loads(sizelistfile.read())
+        except json.decoder.JSONDecodeError:
+            sizelist = {}
+    return sizelist
+
+
+def update_sizelist(sizelist: Dict[str, Dict[str, int]], folder: str) -> None:
+    """
+    Updates the size list JSON file.
+
+    Args:
+        sizelist (Dict[str, Dict[str, int]]): The size list dictionary to be written to the file.
+        folder (str): The folder in which the size list file is located.
+    """
+    sizelist_path = os.path.join(folder, ".sizelist.json")
+    if sizelist != {}:
+        with open(sizelist_path, "w", encoding="utf-8") as sizelistfile:
+            sizelistfile.write(json.dumps(sizelist, indent=4))
+    else:
+        if os.path.exists(sizelist_path):
+            os.remove(sizelist_path)
+
+
+def get_image_info(item: str, folder: str) -> Dict[str, Any]:
+    """
+    Extracts image information and EXIF data.
+
+    Args:
+        item (str): The image file name.
+        folder (str): The folder containing the image.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing image width, height, and EXIF data.
+    """
+    with Image.open(os.path.join(folder, item)) as img:
+        exif = img.getexif()
+        width, height = img.size
+    exifdata = {ExifTags.TAGS.get(key, key): val for key, val in exif.items()}
+    if "Orientation" in exifdata and exifdata["Orientation"] in [6, 8]:
+        width, height = height, width
+    return {"width": width, "height": height}
+
+
+def process_image(item: str, folder: str, _args: Args, baseurl: str, sizelist: Dict[str, Dict[str, int]], raw: List[str]) -> Dict[str, Any]:
+    """
+    Processes an image and prepares its data for the HTML template.
+
+    Args:
+        item (str): The image file name.
+        folder (str): The folder containing the image.
+        _args (Args): Parsed command line arguments.
+        baseurl (str): Base URL for the web root.
+        sizelist (Dict[str, Dict[str, int]]): Dictionary containing size information for images.
+        raw (List[str]): List of raw image file extensions.
+
+    Returns:
+        Dict[str, Any]: Dictionary containing image details for HTML rendering.
+    """
+    extsplit = os.path.splitext(item)
+    if item not in sizelist or _args.regenerate_thumbnails:
+        sizelist[item] = get_image_info(item, folder)
+
+    image = {
+        "url": f"{_args.web_root_url}{baseurl}{urllib.parse.quote(item)}",
+        "thumbnail": f"{_args.web_root_url}.thumbnails/{baseurl}{urllib.parse.quote(item)}.jpg",
+        "name": item,
+        "width": sizelist[item]["width"],
+        "height": sizelist[item]["height"],
+    }
+    if not os.path.exists(os.path.join(_args.root_directory, ".thumbnails", baseurl, item)):
+        thumbnails.append((folder, item, _args.root_directory, _args.regenerate_thumbnails))
+
+    for _raw in raw:
+        if os.path.exists(os.path.join(folder, extsplit[0] + _raw)):
+            url = urllib.parse.quote(extsplit[0]) + _raw
+            if _raw in (".tif", ".tiff"):
+                image["tiff"] = f"{_args.web_root_url}{baseurl}{url}"
+            else:
+                image["raw"] = f"{_args.web_root_url}{baseurl}{url}"
+    return image
+
+
+def generate_html(folder: str, title: str, _args: Args, raw: List[str]) -> None:
+    """
+    Generates HTML content for a folder of images.
+
+    Args:
+        folder (str): The folder to generate HTML for.
+        title (str): The title of the HTML page.
+        _args (Args): Parsed command line arguments.
+        raw (List[str]): Raw image file names.
+    """
+    sizelist = initialize_sizelist(folder)
+    items = sorted(os.listdir(folder))
+
+    images = []
+    subfolders = []
+    foldername = folder.removeprefix(_args.root_directory)
+    foldername = f"{foldername}/" if foldername else ""
+    baseurl = urllib.parse.quote(foldername)
+
+    create_thumbnail_folder(foldername, _args.root_directory)
+
+    if not _args.non_interactive_mode:
+        pbardict[folder] = tqdm(total=len(items), desc=f"Getting image infos - {folder}", unit="files", ascii=True, dynamic_ncols=True)
+
+    for item in items:
+        if item not in EXCLUDES:
+            if os.path.isdir(os.path.join(folder, item)):
+                process_subfolder(item, folder, baseurl, subfolders, _args, raw)
+            else:
+                if os.path.splitext(item)[1].lower() in _args.file_extensions:
+                    images.append(process_image(item, folder, _args, baseurl, sizelist, raw))
+                if item == "info":
+                    process_info_file(folder, item)
+
+        if not _args.non_interactive_mode:
+            pbardict[folder].update(1)
+
+    if not _args.non_interactive_mode:
+        pbardict[folder].close()
+
+    update_sizelist(sizelist, folder)
+
+    if should_generate_html(images, _args):
+        create_html_file(folder, title, foldername, images, subfolders, _args)
+
+    if not _args.non_interactive_mode:
+        pbardict["htmlbar"].update(1)
+
+
+def create_thumbnail_folder(foldername: str, root_directory: str) -> None:
+    """
+    Creates a folder for thumbnails if it doesn't exist.
+
+    Args:
+        foldername (str): The name of the folder.
+        root_directory (str): The root directory path.
+    """
+    thumbnails_path = os.path.join(root_directory, ".thumbnails", foldername)
+    if not os.path.exists(thumbnails_path):
+        os.mkdir(thumbnails_path)
+
+
+def process_subfolder(item: str, folder: str, baseurl: str, subfolders: List[Dict[str, str]], _args: Args, raw: List[str]) -> None:
+    """
+    Processes a subfolder.
+
+    Args:
+        item (str): The name of the subfolder.
+        folder (str): The parent folder containing the subfolder.
+        baseurl (str): Base URL for the web root.
+        subfolders (List[Dict[str, str]]): List to store subfolder details.
+        _args (Args): Parsed command line arguments.
+        raw (List[str]): Raw image file extensions.
+    """
+    subfolder_url = (
+        f"{_args.web_root_url}{baseurl}{urllib.parse.quote(item)}/index.html"
+        if _args.web_root_url.startswith("file://")
+        else f"{_args.web_root_url}{baseurl}{urllib.parse.quote(item)}"
+    )
+    subfolders.append({"url": subfolder_url, "name": item})
+    if item not in _args.exclude_folders:
+        if not any(fnmatch.fnmatchcase(os.path.join(folder, item), exclude) for exclude in _args.exclude_folders):
+            generate_html(os.path.join(folder, item), os.path.join(folder, item).removeprefix(_args.root_directory), _args, raw)
+
+
+def process_info_file(folder: str, item: str) -> None:
+    """
+    Processes an info file.
+
+    Args:
+        folder (str): The folder containing the info file.
+        item (str): The info file name.
+    """
+    with open(os.path.join(folder, item), encoding="utf-8") as f:
+        info[urllib.parse.quote(folder)] = f.read()
+
+
+def should_generate_html(images: List[Dict[str, Any]], _args: Args) -> bool:
+    """
+    Determines if HTML should be generated.
+
+    Args:
+        images (List[Dict[str, Any]]): List of images.
+        _args (Args): Parsed command line arguments.
+
+    Returns:
+        bool: True if HTML should be generated, False otherwise.
+    """
+    return images or (_args.use_fancy_folders and (not images or _args.ignore_other_files))
+
+
+def create_html_file(
+    folder: str, title: str, foldername: str, images: List[Dict[str, Any]], subfolders: List[Dict[str, str]], _args: Args
+) -> None:
+    """
+    Creates the HTML file using the template.
+
+    Args:
+        folder (str): The folder to create the HTML file in.
+        title (str): The title of the HTML page.
+        foldername (str): The name of the folder.
+        images (List[Dict[str, Any]]): A list of images to include in the HTML.
+        subfolders (List[Dict[str, str]]): A list of subfolders to include in the HTML.
+        _args (Args): Parsed command line arguments.
+    """
+    image_chunks = np.array_split(images, 8) if images else []
+    header = os.path.basename(folder) or title
+    parent = None if not foldername else f"{_args.web_root_url}{urllib.parse.quote(foldername.removesuffix(folder.split('/')[-1] + '/'))}"
+    if parent and _args.web_root_url.startswith("file://"):
+        parent += "index.html"
+
+    license_info = (
+        {
+            "project": _args.site_title,
+            "author": _args.author_name,
+            "type": cclicense.licensenameswitch(_args.license_type),
+            "url": cclicense.licenseurlswitch(_args.license_type),
+            "pics": cclicense.licensepicswitch(_args.license_type),
+        }
+        if _args.license_type
+        else None
+    )
+
+    folder_info = info.get(urllib.parse.quote(folder), "").split("\n")
+    _info = [i for i in folder_info if len(i) > 1] if folder_info else None
+
+    html = env.get_template("index.html.j2")
+    content = html.render(
+        title=title,
+        favicon=f"{_args.web_root_url}{FAVICON_PATH}",
+        stylesheet=f"{_args.web_root_url}{GLOBAL_CSS_PATH}",
+        theme=f"{_args.web_root_url}.static/theme.css",
+        root=_args.web_root_url,
+        parent=parent,
+        header=header,
+        license=license_info,
+        subdirectories=subfolders,
+        images=image_chunks,
+        info=_info,
+        allimages=images,
+        webmanifest=_args.generate_webmanifest,
+    )
+
+    with open(os.path.join(folder, "index.html"), "w", encoding="utf-8") as f:
+        f.write(content)
+
+
+def list_folder(total: int, folder: str, title: str, _args: Args, raw: List[str]) -> List[Tuple[str, str]]:
+    """
+    Lists and processes a folder, generating HTML files.
+
+    Args:
+        total (int): Total number of folders to process.
+        folder (str): The folder to process.
+        title (str): The title of the HTML page.
+        _args (Args): Parsed command line arguments.
+        raw (List[str]): Raw image file names.
+
+    Returns:
+        List[Tuple[str, str]]: List of thumbnails generated.
+    """
+    if not _args.non_interactive_mode:
+        pbardict["htmlbar"] = tqdm(total=total, desc="Generating HTML files", unit="folders", ascii=True, dynamic_ncols=True)
+    generate_html(folder, title, _args, raw)
+    return thumbnails