added tagging function as requested in issue #7

modules/generate_html.py

@@ -7,7 +7,7 @@ from typing import Any
 from datetime import datetime
 
 from tqdm.auto import tqdm
-from PIL import Image, ExifTags, TiffImagePlugin
+from PIL import Image, ExifTags, TiffImagePlugin, UnidentifiedImageError
 from jinja2 import Environment, FileSystemLoader
 
 from modules.logger import logger
@@ -29,54 +29,60 @@ Image.MAX_IMAGE_PIXELS = 933120000
 
 # Initialize Jinja2 environment for template rendering
 env = Environment(loader=FileSystemLoader(os.path.join(SCRIPTDIR, "templates")))
-thumbnails: list[tuple[str, str]] = []
+thumbnails: list[tuple[str, str, str]] = []
 info: dict[str, str] = {}
 licens: dict[str, str] = {}
 
 
-def initialize_sizelist(folder: str) -> dict[str, dict[str, int]]:
+def initialize_metadata(folder: str) -> dict[str, dict[str, int]]:
     """
-    Initializes the size list JSON file if it doesn't exist.
+    Initializes the metadata JSON file if it doesn't exist.
 
     Args:
-        folder (str): The folder in which the size list file is located.
+        folder (str): The folder in which the metadata file is located.
 
     Returns:
-        dict[str, dict[str, int]]: The size list dictionary.
+        dict[str, dict[str, int]]: The metadata dictionary.
     """
-    sizelist = {}
-    sizelist_path = os.path.join(folder, ".sizelist.json")
-    if not os.path.exists(sizelist_path):
-        logger.info("creating new size list file", extra={"file": sizelist_path})
-        with open(sizelist_path, "x", encoding="utf-8") as sizelistfile:
-            sizelistfile.write("{}")
-    with open(sizelist_path, "r+", encoding="utf-8") as sizelistfile:
-        logger.info("reading size list file", extra={"file": sizelist_path})
+    metadata = {}
+    metadata_path = os.path.join(folder, ".metadata.json")
+    if not os.path.exists(metadata_path):
+        logger.info("creating new metadata file", extra={"file": metadata_path})
+        with open(metadata_path, "x", encoding="utf-8") as metadatafile:
+            metadatafile.write("{}")
+    with open(metadata_path, "r+", encoding="utf-8") as metadatafile:
+        logger.info("reading metadata file", extra={"file": metadata_path})
         try:
-            sizelist = json.loads(sizelistfile.read())
+            metadata = json.loads(metadatafile.read())
         except json.decoder.JSONDecodeError:
-            logger.warning("invalid JSON in size list file", extra={"file": sizelist_path})
-            sizelist = {}
-    return sizelist
+            logger.warning("invalid JSON in metadata file", extra={"file": metadata_path})
+            metadata = {}
+
+    # remove old sizelist if it exists
+    sizelist_path = os.path.join(folder, ".sizelist.json")
+    if os.path.exists(sizelist_path):
+        logger.warning("found old .sizelist.json, removing it...", extra={"path": sizelist_path})
+        os.remove(sizelist_path)
+    return metadata
 
 
-def update_sizelist(sizelist: dict[str, dict[str, Any]], folder: str) -> None:
+def update_metadata(metadata: dict[str, dict[str, Any]], folder: str) -> None:
     """
-    Updates the size list JSON file.
+    Updates the metadata 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.
+        metadata (dict[str, dict[str, int]]): The metadata dictionary to be written to the file.
+        folder (str): The folder in which the metadata file is located.
     """
-    sizelist_path = os.path.join(folder, ".sizelist.json")
-    if sizelist:
-        with open(sizelist_path, "w", encoding="utf-8") as sizelistfile:
-            logger.info("writing size list file", extra={"file": sizelist_path})
-            sizelistfile.write(json.dumps(sizelist, indent=4))
+    metadata_path = os.path.join(folder, ".metadata.json")
+    if metadata:
+        with open(metadata_path, "w", encoding="utf-8") as metadatafile:
+            logger.info("writing metadata file", extra={"file": metadata_path})
+            metadatafile.write(json.dumps(metadata, indent=4))
     else:
-        if os.path.exists(sizelist_path):
-            logger.info("deleting empty size list file", extra={"file": sizelist_path})
-            os.remove(sizelist_path)
+        if os.path.exists(metadata_path):
+            logger.info("deleting empty metadata file", extra={"file": metadata_path})
+            os.remove(metadata_path)
 
 
 def get_image_info(item: str, folder: str) -> dict[str, Any]:
@@ -91,10 +97,15 @@ def get_image_info(item: str, folder: str) -> dict[str, Any]:
         dict[str, Any]: A dictionary containing image width, height, and EXIF data.
     """
     file = os.path.join(folder, item)
-    with Image.open(file) as img:
-        logger.info("extracting image information", extra={"file": file})
-        exif = img.getexif()
-        width, height = img.size
+    try:
+        with Image.open(file) as img:
+            logger.info("extracting image information", extra={"file": file})
+            exif = img.getexif()
+            width, height = img.size
+    except UnidentifiedImageError:
+        logger.error("cannot identify image file", extra={"file": file})
+        print(f"cannot identify image file: {file}")
+        return {"width": None, "height": None, "tags": None, "exifdata": None}
     if exif:
         logger.info("extracting EXIF data", extra={"file": file})
         ifd = exif.get_ifd(ExifTags.IFD.Exif)
@@ -116,9 +127,9 @@ def get_image_info(item: str, folder: str) -> dict[str, Any]:
                     content = newtuple
             if tag in ["DateTime", "DateTimeOriginal", "DateTimeDigitized"]:
                 epr = r"\d{4}:\d{2}:\d{2} \d{2}:\d{2}:\d{2}"
-                if re.match(epr, content):
+                if re.match(epr, str(content)):
                     try:
-                        content = datetime.strptime(content, "%Y:%m:%d %H:%M:%S").strftime("%Y-%m-%d %H:%M:%S")
+                        content = datetime.strptime(str(content), "%Y:%m:%d %H:%M:%S").strftime("%Y-%m-%d %H:%M:%S")
                     except ValueError:
                         content = None
                 else:
@@ -130,12 +141,12 @@ def get_image_info(item: str, folder: str) -> dict[str, Any]:
         for key in ["PrintImageMatching", "UserComment", "MakerNote"]:
             if key in exifdata:
                 del exifdata[key]
-        return {"width": width, "height": height, "exifdata": exifdata}
+        return {"width": width, "height": height, "tags": [], "exifdata": exifdata}
     else:
-        return {"width": width, "height": height, "exifdata": None}
+        return {"width": width, "height": height, "tags": [], "exifdata": None}
 
 
-def process_image(item: str, folder: str, _args: Args, baseurl: str, sizelist: dict[str, dict[str, int]], raw: list[str]) -> dict[str, Any]:
+def process_image(item: str, folder: str, _args: Args, baseurl: str, metadata: dict[str, dict[str, int]], raw: list[str]) -> dict[str, Any]:
     """
     Processes an image and prepares its data for the HTML template.
 
@@ -144,23 +155,24 @@ def process_image(item: str, folder: str, _args: Args, baseurl: str, sizelist: d
         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.
+        metadata (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.reread_metadata:
-        sizelist[item] = get_image_info(item, folder)
+    if item not in metadata or _args.reread_metadata:
+        metadata[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"],
-        "exifdata": sizelist[item].get("exifdata", ""),
+        "width": metadata[item]["width"],
+        "height": metadata[item]["height"],
+        "tags": metadata[item]["tags"],
+        "exifdata": metadata[item].get("exifdata", ""),
     }
     path = os.path.join(_args.root_directory, ".thumbnails", baseurl, item + ".jpg")
     if not os.path.exists(path) or _args.regenerate_thumbnails:
@@ -193,10 +205,10 @@ def generate_html(folder: str, title: str, _args: Args, raw: list[str], version:
     """
     logger.info("processing folder", extra={"folder": folder})
     if _args.regenerate_thumbnails:
-        if os.path.exists(os.path.join(folder, ".sizelist.json")):
-            logger.info("removing .sizelist.json", extra={"folder": folder})
-            os.remove(os.path.join(folder, ".sizelist.json"))
-    sizelist = initialize_sizelist(folder)
+        if os.path.exists(os.path.join(folder, ".metadata.json")):
+            logger.info("removing .metadata.json", extra={"folder": folder})
+            os.remove(os.path.join(folder, ".metadata.json"))
+    metadata = initialize_metadata(folder)
     items = sorted(os.listdir(folder))
 
     contains_files = False
@@ -217,7 +229,7 @@ def generate_html(folder: str, title: str, _args: Args, raw: list[str], version:
                 else:
                     contains_files = True
                     if os.path.splitext(item)[1].lower() in _args.file_extensions:
-                        images.append(process_image(item, folder, _args, baseurl, sizelist, raw))
+                        images.append(process_image(item, folder, _args, baseurl, metadata, raw))
                     if item == "info":
                         process_info_file(folder, item)
                     if item == "LICENSE":
@@ -230,13 +242,13 @@ def generate_html(folder: str, title: str, _args: Args, raw: list[str], version:
                 else:
                     contains_files = True
                     if os.path.splitext(item)[1].lower() in _args.file_extensions:
-                        images.append(process_image(item, folder, _args, baseurl, sizelist, raw))
+                        images.append(process_image(item, folder, _args, baseurl, metadata, raw))
                     if item == "info":
                         process_info_file(folder, item)
                     if item == "LICENSE":
                         process_license(folder, item)
 
-    update_sizelist(sizelist, folder)
+    update_metadata(metadata, folder)
 
     if should_generate_html(images, contains_files, _args):
         create_html_file(folder, title, foldername, images, subfolders, _args, version, logo)
@@ -260,7 +272,7 @@ def create_thumbnail_folder(foldername: str, root_directory: str) -> None:
         os.mkdir(thumbnails_path)
 
 
-def process_subfolder(item: str, folder: str, baseurl: str, subfolders: list[dict[str, str]], _args: Args, raw: list[str], version: str, logo: str) -> None:
+def process_subfolder(item: str, folder: str, baseurl: str, subfolders: list[dict[str, str | None]], _args: Args, raw: list[str], version: str, logo: str) -> None:
     """
     Processes a subfolder.
 
@@ -365,6 +377,12 @@ def create_html_file(folder: str, title: str, foldername: str, images: list[dict
         else None
     )
 
+    alltags = set()
+    for img in images:
+        for tag in img["tags"]:
+            alltags.add(tag)
+    alltags = sorted(alltags)
+
     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
     if _args.reverse_sort:
@@ -413,6 +431,7 @@ def create_html_file(folder: str, title: str, foldername: str, images: list[dict
         version=version,
         logo=logo,
         licensefile=license_url,
+        tags=alltags,
     )
 
     with open(html_file, "w", encoding="utf-8") as f:
@@ -420,7 +439,7 @@ def create_html_file(folder: str, title: str, foldername: str, images: list[dict
         f.write(content)
 
 
-def list_folder(folder: str, title: str, _args: Args, raw: list[str], version: str, logo: str) -> list[tuple[str, str]]:
+def list_folder(folder: str, title: str, _args: Args, raw: list[str], version: str, logo: str) -> list[tuple[str, str, str]]:
     """
     lists and processes a folder, generating HTML files.