dev(narugo): add more docs

imgutils/metadata/geninfo.py

@@ -1,3 +1,16 @@
+"""
+This module provides functions for reading and writing generation information (geninfo) to image files.
+It supports different image formats including PNG, EXIF, and GIF.
+
+The module includes functions for:
+
+1. Reading geninfo from image parameters, EXIF data, and GIF comments
+2. Writing geninfo to image parameters, EXIF data, and GIF comments
+
+These functions are useful for storing and retrieving metadata about image generation,
+particularly in the context of AI-generated images.
+"""
+
 from typing import Optional
 
 import piexif
@@ -8,12 +21,38 @@
 
 
 def read_geninfo_parameters(image: ImageTyping) -> Optional[str]:
+    """
+    Read generation information from image parameters.
+
+    :param image: The input image.
+    :type image: ImageTyping
+
+    :return: The generation information if found, None otherwise.
+    :rtype: Optional[str]
+
+    This function loads the image and attempts to retrieve the 'parameters' 
+    information from the image metadata. It's commonly used for PNG images 
+    where generation information is stored in the image parameters.
+    """
     image = load_image(image, mode=None, force_background=None)
     infos = image.info or {}
     return infos.get('parameters')
 
 
 def read_geninfo_exif(image: ImageTyping) -> Optional[str]:
+    """
+    Read generation information from EXIF data.
+
+    :param image: The input image.
+    :type image: ImageTyping
+
+    :return: The generation information if found in EXIF data, None otherwise.
+    :rtype: Optional[str]
+
+    This function attempts to read generation information from the EXIF metadata
+    of the image. It specifically looks for the UserComment field in the EXIF data.
+    If the EXIF data is invalid or not present, it returns None.
+    """
     image = load_image(image, mode=None, force_background=None)
     infos = image.info or {}
     if "exif" in infos:
@@ -36,6 +75,19 @@ def read_geninfo_exif(image: ImageTyping) -> Optional[str]:
 
 
 def read_geninfo_gif(image: ImageTyping) -> Optional[str]:
+    """
+    Read generation information from GIF comment.
+
+    :param image: The input image.
+    :type image: ImageTyping
+
+    :return: The generation information if found in GIF comment, None otherwise.
+    :rtype: Optional[str]
+
+    This function is specifically designed to read generation information from
+    GIF images. It looks for the 'comment' field in the image metadata, which is
+    commonly used in GIF files to store additional information.
+    """
     image = load_image(image, mode=None, force_background=None)
     infos = image.info or {}
     if "comment" in infos:  # for gif
@@ -45,23 +97,69 @@ def read_geninfo_gif(image: ImageTyping) -> Optional[str]:
 
 
 def write_geninfo_parameters(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs):
+    """
+    Write generation information to image parameters.
+
+    :param image: The input image.
+    :type image: ImageTyping
+    :param dst_filename: The destination filename to save the image with geninfo.
+    :type dst_filename: str
+    :param geninfo: The generation information to write.
+    :type geninfo: str
+    :param kwargs: Additional keyword arguments to pass to the image save function.
+
+    This function writes the provided generation information to the image parameters.
+    It's commonly used for PNG images where generation information can be stored in
+    the image metadata. The function creates a PngInfo object, adds the geninfo as
+    'parameters', and saves the image with this metadata.
+    """
     pnginfo = PngInfo()
     pnginfo.add_text('parameters', geninfo)
 
     image = load_image(image, force_background=None, mode=None)
-    image.save(dst_filename, pnginfo=pnginfo, *kwargs)
+    image.save(dst_filename, pnginfo=pnginfo, **kwargs)
 
 
 def write_geninfo_exif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs):
+    """
+    Write generation information to EXIF data.
+
+    :param image: The input image.
+    :type image: ImageTyping
+    :param dst_filename: The destination filename to save the image with geninfo.
+    :type dst_filename: str
+    :param geninfo: The generation information to write.
+    :type geninfo: str
+    :param kwargs: Additional keyword arguments to pass to the image save function.
+
+    This function writes the provided generation information to the EXIF metadata
+    of the image. It creates an EXIF dictionary with the geninfo stored in the
+    UserComment field, converts it to bytes, and saves the image with this EXIF data.
+    """
     exif_dict = {
         "Exif": {piexif.ExifIFD.UserComment: UserComment.dump(geninfo, encoding="unicode")}}
     exif_bytes = piexif.dump(exif_dict)
 
     image = load_image(image, force_background=None, mode=None)
-    image.save(dst_filename, exif=exif_bytes, *kwargs)
+    image.save(dst_filename, exif=exif_bytes, **kwargs)
 
 
 def write_geninfo_gif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs):
+    """
+    Write generation information to GIF comment.
+
+    :param image: The input image.
+    :type image: ImageTyping
+    :param dst_filename: The destination filename to save the image with geninfo.
+    :type dst_filename: str
+    :param geninfo: The generation information to write.
+    :type geninfo: str
+    :param kwargs: Additional keyword arguments to pass to the image save function.
+
+    This function is specifically designed to write generation information to
+    GIF images. It adds the geninfo to the image's 'comment' field, which is
+    a standard way of including metadata in GIF files.
+    """
     image = load_image(image, force_background=None, mode=None)
     image.info['comment'] = geninfo.encode('utf-8')
-    image.save(dst_filename, *kwargs)
+    image.save(dst_filename, **kwargs)

imgutils/sd/nai.py

@@ -59,6 +59,12 @@ class NAIMetadata:
 
     @property
     def json(self) -> dict:
+        """
+        Convert the NAIMetadata to a JSON-compatible dictionary.
+
+        :return: A dictionary representation of the metadata.
+        :rtype: dict
+        """
         data = {
             'Software': self.software,
             'Source': self.source,
@@ -90,10 +96,22 @@ def pnginfo(self) -> PngInfo:
 
 
 class _InvalidNAIMetaError(Exception):
+    """
+    Custom exception raised when NAI metadata is invalid.
+    """
     pass
 
 
 def _naimeta_validate(data):
+    """
+    Validate the NAI metadata.
+
+    :param data: The metadata to validate.
+    :type data: dict
+    :return: The validated metadata.
+    :rtype: dict
+    :raises _InvalidNAIMetaError: If the metadata is invalid.
+    """
     if isinstance(data, dict) and data.get('Software') and data.get('Source') and data.get('Comment'):
         return data
     else:
@@ -104,14 +122,14 @@ def _get_naimeta_raw(image: ImageTyping) -> dict:
     """
     Extract raw NAI metadata from an image.
 
-    This function attempts to extract metadata from the image using LSB (Least Significant Bit) extraction.
-    If that fails, it falls back to using the image's info dictionary.
+    This function attempts to extract metadata from the image using various methods,
+    including LSB extraction, image info dictionary, and other specific metadata formats.
 
     :param image: The input image.
     :type image: ImageTyping
-
     :return: A dictionary containing the raw metadata.
     :rtype: dict
+    :raises _InvalidNAIMetaError: If no valid metadata is found.
     """
     image = load_image(image, force_background=None, mode=None)
     try:
@@ -149,7 +167,6 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
 
     :param image: The input image.
     :type image: ImageTyping
-
     :return: A NAIMetadata object if successful, None otherwise.
     :rtype: Optional[NAIMetadata]
     """
@@ -169,19 +186,62 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
 
 
 def add_naimeta_to_image(image: ImageTyping, metadata: NAIMetadata) -> Image.Image:
+    """
+    Add NAI metadata to an image using LSB (Least Significant Bit) encoding.
+
+    :param image: The input image.
+    :type image: ImageTyping
+    :param metadata: The NAIMetadata object to add to the image.
+    :type metadata: NAIMetadata
+    :return: The image with added metadata.
+    :rtype: Image.Image
+    """
     image = load_image(image, mode=None, force_background=None)
     return write_lsb_metadata(image, data=metadata.pnginfo)
 
 
 def _save_png_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):
+    """
+    Save a PNG image with NAI metadata.
+
+    :param image: The image to save.
+    :type image: Image.Image
+    :param dst_file: The destination file path.
+    :type dst_file: Union[str, os.PathLike]
+    :param metadata: The NAIMetadata object to include in the image.
+    :type metadata: NAIMetadata
+    :param kwargs: Additional keyword arguments for image saving.
+    """
     image.save(dst_file, pnginfo=metadata.pnginfo, **kwargs)
 
 
 def _save_exif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):
+    """
+    Save an image with NAI metadata in EXIF format.
+
+    :param image: The image to save.
+    :type image: Image.Image
+    :param dst_file: The destination file path.
+    :type dst_file: Union[str, os.PathLike]
+    :param metadata: The NAIMetadata object to include in the image.
+    :type metadata: NAIMetadata
+    :param kwargs: Additional keyword arguments for image saving.
+    """
     write_geninfo_exif(image, dst_file, json.dumps(metadata.json), **kwargs)
 
 
 def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):
+    """
+    Save a GIF image with NAI metadata.
+
+    :param image: The image to save.
+    :type image: Image.Image
+    :param dst_file: The destination file path.
+    :type dst_file: Union[str, os.PathLike]
+    :param metadata: The NAIMetadata object to include in the image.
+    :type metadata: NAIMetadata
+    :param kwargs: Additional keyword arguments for image saving.
+    """
     write_geninfo_gif(image, dst_file, json.dumps(metadata.json), **kwargs)
 
 
@@ -197,6 +257,28 @@ def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike]
 def save_image_with_naimeta(
         image: ImageTyping, dst_file: Union[str, os.PathLike], metadata: NAIMetadata,
         add_lsb_meta: Union[str, bool] = 'auto', save_metainfo: Union[str, bool] = 'auto', **kwargs) -> Image.Image:
+    """
+    Save an image with NAI metadata.
+
+    This function saves the given image with the provided NAI metadata. It can add LSB metadata
+    and save metainfo based on the image format and user preferences.
+
+    :param image: The input image.
+    :type image: ImageTyping
+    :param dst_file: The destination file path.
+    :type dst_file: Union[str, os.PathLike]
+    :param metadata: The NAIMetadata object to include in the image.
+    :type metadata: NAIMetadata
+    :param add_lsb_meta: Whether to add LSB metadata. Can be 'auto', True, or False.
+    :type add_lsb_meta: Union[str, bool]
+    :param save_metainfo: Whether to save metainfo. Can be 'auto', True, or False.
+    :type save_metainfo: Union[str, bool]
+    :param kwargs: Additional keyword arguments for image saving.
+    :return: The saved image.
+    :rtype: Image.Image
+    :raises ValueError: If LSB metadata cannot be saved to the specified image format.
+    :raises SystemError: If the image format is not supported for saving metainfo.
+    """
     mimetype, _ = mimetypes.guess_type(dst_file)
     if add_lsb_meta == 'auto':
         if mimetype in _LSB_ALLOWED_TYPES: