mee/mee.py

"""Minimalistic EXIF Editor.

Commands overview:
    ls, get, set, del       operate with EXIF tags.
    cget, cset, cdel        operate with JSON in UserComment EXIF tag.

Usage:
    mee ls [options] <file>
    mee get [options] [-k <key>] <file>
    mee set [options] -k <key> -v <value> <file> [-o <file>]
    mee del [options] (--all | -k <key>) <file> [-o <file>]
    mee cget [options] [-k <key>] <file>
    mee cset [options] -k <key> -v <value> <file> [-o <file>]
    mee cdel [options] -k <key> <file> [-o <file>]
    mee (--help | --version)

Options:
    -o, --output <file>     output file. Same as input file by default.
    -k, --key <key>         set EXIF/JSON key.
    -v, --value <value>     EXIF/JSON key value.
    -c, --custom            make 'get -c' alias for 'cget', etc.
    --json                  print JSON output if available.
    --all                   perform action with all EXIF tags.
    --debug                 enable debug messages.
    --help                  print this message and exit.
    --version               print version and exit.

Environment:
    DEBUG                   enable debug messages.
"""

__version__ = "0.2.1"

import os
import sys
import json
import logging
from tempfile import NamedTemporaryFile
from typing import Any, Optional
from pathlib import Path

from exif import Image
from docopt import docopt


class Mee:
    """
    Minimalistic EXIF Editor class. Custom tags implemented as JSON string
    written in Exif.Photo.UserComment tag.

    Properties:
        infile: `Path`          Path to input file.
        storage: `str`          EXIF tag name used as custom tags storage.
        image: `exif.Image`     exif.Image object.
        tags: `dict`            Custom tags.

    Methods:
        get()       Get custom tag value or all tags.
        set()       Set new custom tag or modify existing.
        delete()    Remove custom tag.
        commit()    Write changs into file on disk.

    You can directly use `exif.Image` methods::

        mee = Mee("image.jpg")
        mee.image.get_all()  # get all EXIF tags

    See also:
    * https://gitlab.com/TNThieding/exif/-/blob/master/src/exif/_image.py
    * https://exif.readthedocs.io/en/latest/usage.html
    * https://exiv2.org/doc/exifcomment_8cpp-example.html
    """

    def __init__(self, infile: Path):
        self.infile = infile
        self.storage = "user_comment"
        self.logger = logging.getLogger("Mee")

        with open(infile, "rb") as in_img:
            self.image = Image(in_img)

        try:
            tags = self.exif_get(self.storage) or "{}"
            self.logger.debug("Raw custom tags: %s", tags)
            if tags:
                self.tags = json.loads(tags)
                try:
                    # See WARNING in commit() method comments.
                    del self.tags["_"]
                except KeyError:
                    pass
                self.logger.debug("Got self.tags from JSON: %s", self.tags)
        except (KeyError, json.JSONDecodeError) as error:
            self.tags = {}
            self.logger.debug(
                "Set empty dict self.tags on error: %s", self.tags, error
            )

    def exif_get(self, key: str) -> str:
        """Return EXIF tag value."""
        self.logger.debug("Call exif_get(): key=%s", key)
        try:
            return self.image.get(key, None)
        except KeyError:
            return None

    def exif_set(self, key: str, value: Any) -> None:
        """Set EXIF tag. Wraps `exif.Image.set()`."""
        self.logger.debug("Call exif_set(): key=%s, value=%s", key, value)
        self.image.set(key, value)

    def exif_del(self, key: str) -> None:
        """Delete EXIF tag. Wraps `exif.Image.delete()`."""
        self.logger.debug("Call exif_del(): key=%s", key)
        try:
            self.image.delete(key)
        except (KeyError, AttributeError) as error:
            self.logger.debug("Pass on error in exif_del(): %s", error)

    def get(self, key: str = None) -> Any:
        """Return custom tag value or all custom tags."""
        if key:
            try:
                return self.tags[key]
            except KeyError:
                return None
        return self.tags

    def set(self, key: str, value: Any) -> None:
        """Set new or edit existing custom tag key:value pair."""
        self.logger.debug("Call set(): key=%s, value=%s", key, value)
        if not self.tags:
            self.tags = {}
        self.tags[key] = value

    def delete(self, key: str) -> None:
        """Remove custom tag."""
        self.logger.debug("Call delete(): key=%s", key)
        if self.tags:
            try:
                self.logger.debug("Call delete(): delete key")
                del self.tags[key]
            except KeyError:
                self.logger.debug("Pass on KeyError in delete()")

    def write_out(self, outfile: Optional[Path] = None) -> None:
        """Just write data to `outfile` or overwrite `infile`."""
        if not outfile:
            with NamedTemporaryFile() as out_tmp, open(
                self.infile, "wb"
            ) as out_img:
                out_tmp.write(self.image.get_file())
                out_tmp.seek(0)
                out_img.write(out_tmp.read())
        else:
            with open(outfile, "wb") as out_img:
                out_img.write(self.image.get_file())

    def commit(self, outfile: Optional[Path] = None) -> None:
        """Safely write changes to file."""
        self.logger.debug("Call commit(): outfile=%s", outfile)
        # Remove old tag from EXIF.
        # You cannot replace the data if it is longer than it was originally.
        # Therefore, old data from the tag is always preliminarily removed.
        self.exif_del(self.storage)
        if not self.tags:
            # !!! WARNING !!!
            # If you write empty JSON to EXIF, then EXIF will be considered
            # corrupted! You'll get a warning in exiftool, and you won't be
            # able to do anything with the metadata. Here I add a placeholder
            # key so that valid JSON is always written to UserComment tag.
            self.tags = {"_": None}
        string = json.dumps(self.tags, separators=(",", ":"))
        self.logger.debug("Commit: self.tags: %s", self.tags)
        self.logger.debug("Commit: Data to write: %s", string)
        self.exif_set(self.storage, string)
        # Write data to file
        self.write_out(outfile)


def cli() -> None:
    args = docopt(__doc__, version=__version__)
    key = args["--key"]
    value = args["--value"]
    infile = args["<file>"]
    outfile = args["--output"]

    if args["--debug"] or os.getenv("DEBUG"):
        logging.basicConfig(level=logging.DEBUG)

    mee = Mee(infile)

    if args["ls"]:
        print(mee.image.list_all())

    if args["get"] and not args["--custom"]:
        if not key:
            tags = mee.image.get_all()
            if args['--json']:
                print(json.dumps(tags, indent=4))
            else:
                print(tags)
        else:
            tag = mee.exif_get(key)
            if args['--json']:
                print(json.dumps(tag, indent=4))
            else:
                print(tag)

    if args["set"] and not args["--custom"]:
        mee.exif_del(mee.storage)
        mee.exif_set(key, value)
        mee.write_out(outfile)

    if args["del"] and not args["--custom"]:
        if args["--all"]:
            mee.image.delete_all()
        else:
            mee.image.delete(key)
        mee.write_out(outfile)

    if args["cget"] or (args["get"] and args["--custom"]):
        if not key:
            print(json.dumps(mee.tags, indent=4))
        else:
            try:
                print(mee.tags[key])
            except KeyError:
                sys.exit(1)

    if args["cset"] or (args["set"] and args["--custom"]):
        mee.set(key, value)
        mee.commit(outfile)

    if args["cdel"] or args["del"] and args["--custom"]:
        mee.delete(key)
        mee.commit(outfile)


if __name__ == "__main__":
    cli()