imgcat.py

#!/usr/bin/env python3

"""
Read terminal properties and imgcat in Python.

References:
    https://en.wikipedia.org/wiki/ANSI_escape_code
    https://iterm2.com/documentation-escape-codes.html
    https://iterm2.com/documentation-images.html
"""


import sys
import termios
import re
import binascii
import io


def _send_command(cmd: str) -> str:
    """
    Send a command to the terminal and read the response with 200ms.
    """
    old_stdin_mode = termios.tcgetattr(sys.stdin)
    rep = termios.tcgetattr(sys.stdin)
    rep[3] = rep[3] & ~(termios.ECHO | termios.ICANON)  # lflags
    rep[6][termios.VMIN] = 0  # cc, Minimum number of characters for noncanonical read (MIN).
    rep[6][termios.VTIME] = 2  # cc, Timeout in deciseconds for noncanonical read (TIME).
    termios.tcsetattr(sys.stdin, termios.TCSAFLUSH, rep)
    try:
        sys.stdout.write(cmd)
        sys.stdout.flush()
        return sys.stdin.read()
    finally:
        termios.tcsetattr(sys.stdin, termios.TCSAFLUSH, old_stdin_mode)


def name_version():
    """
    Report terminal name and version.

    Sequence: CSI > q
    """
    cmd = "\033[>q"
    rep = _send_command(cmd)
    rep = re.match("\033P>\\|(.+)(\007|\033\\\\)", rep)
    if rep:
        return rep.group(1)


def cell_size():
    """
    Report the size in points of a single character cell.

    Sequence: OSC 1337 ; ReportCellSize ST
    """
    cmd = "\033]1337;ReportCellSize\007"
    rep = _send_command(cmd)
    rep = re.match("\033\\]1337;ReportCellSize=(.*)(\007|\033\\\\)", rep)
    if rep:
        return rep.group(1).split(";")


def cursor_position():
    """
    Device Status Report (DSR): report the cursor position (CPR)

    Sequence: CSI 6n
    """
    cmd = "\x1b[6n"
    rep = _send_command(cmd)
    rep = re.match("\033\\[(.*)R", rep)
    if rep:
        return rep.group(1).split(";")


def imgcat(img: bytes, name=None, width=None, height=None, preserve_aspect_ratio=True):
    """
    Display an image in a capable terminal (iTerm2, WezTerm, etc.).
    """
    sys.stdout.write(f"\033]1337;File=inline=1;size={len(img)}")
    if name:
        sys.stdout.write(f";name={binascii.b2a_base64(name.encode()).decode()}")
    if width:
        sys.stdout.write(f";width={width}")
    if height:
        sys.stdout.write(f";height={height}")
    if not preserve_aspect_ratio:
        sys.stdout.write(f";preserveAspectRatio=0")
    sys.stdout.write(":")
    sys.stdout.write(binascii.b2a_base64(img).decode())
    sys.stdout.write("\007\n")


class ImgCat(io.BytesIO):
    """
    Class version of the imgcat function.

    Example:
        from PIL import Image
        im = Image.open("image.png")
        im.save(imgcat.ImgCat(height=25), "PNG")
    """

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        self.done = False
        super().__init__()

    def flush(self):
        self.imgcat()
        return super().flush()

    def close(self):
        self.imgcat()
        return super().close()

    def imgcat(self):
        if self.done:
            return
        self.done = True

        if not sys.stdout.isatty():
            return
        img = self.getvalue()
        if len(img) > 0:
            imgcat(img, *self.args, **self.kwargs)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser("Display an image in the terminal")
    parser.add_argument("-H", "--height", metavar="N", type=str)
    parser.add_argument("-W", "--width", metavar="N", type=str)
    parser.add_argument("--no-preserve-aspect-ratio", action="store_true")
    parser.add_argument("files", nargs="*")
    args = parser.parse_args()
    print(args)

    for file in args.files:
        imgcat(
            open(file, "rb").read(),
            width=args.width,
            height=args.height,
            preserve_aspect_ratio=not args.no_preserve_aspect_ratio,
        )