kitty/kitty/icat.py

285 lines
8.9 KiB
Python
Executable File

#!/usr/bin/env python3
# vim:fileencoding=utf-8
# License: GPL v3 Copyright: 2017, Kovid Goyal <kovid at kovidgoyal.net>
import codecs
import fcntl
import mimetypes
import os
import re
import selectors
import signal
import struct
import subprocess
import sys
import termios
import tty
import zlib
from base64 import standard_b64encode
from collections import namedtuple
from math import ceil, floor
from tempfile import NamedTemporaryFile
from time import monotonic
from kitty.constants import appname
from kitty.cli import parse_args
try:
fsenc = sys.getfilesystemencoding() or 'utf-8'
codecs.lookup(fsenc)
except Exception:
fsenc = 'utf-8'
class OpenFailed(ValueError):
def __init__(self, path, message):
ValueError.__init__(
self, 'Failed to open: {} with error: {}'.format(path, message)
)
self.path = path
OPTIONS = '''\
--align
type=choices
choices=center,left,right
default=center
Horizontal alignment for the displayed image. Defaults to centered.
'''
def options_spec():
if not hasattr(options_spec, 'ans'):
options_spec.ans = OPTIONS.format(
appname='{}-icat'.format(appname),
)
return options_spec.ans
Size = namedtuple('Size', 'rows cols width height')
def screen_size(refresh=False):
if refresh or getattr(screen_size, 'ans', None) is None:
s = struct.pack('HHHH', 0, 0, 0, 0)
x = fcntl.ioctl(1, termios.TIOCGWINSZ, s)
screen_size.ans = Size(*struct.unpack('HHHH', x))
return screen_size.ans
def write_gr_cmd(cmd, payload):
cmd = ','.join('{}={}'.format(k, v) for k, v in cmd.items())
w = sys.stdout.buffer.write
w(b'\033_G'), w(cmd.encode('ascii')), w(b';'), w(payload), w(b'\033\\')
sys.stdout.flush()
def fit_image(width, height, pwidth, pheight):
if height > pheight:
corrf = pheight / float(height)
width, height = floor(corrf * width), pheight
if width > pwidth:
corrf = pwidth / float(width)
width, height = pwidth, floor(corrf * height)
if height > pheight:
corrf = pheight / float(height)
width, height = floor(corrf * width), pheight
return int(width), int(height)
def set_cursor(cmd, width, height, align):
ss = screen_size()
cw = int(ss.width / ss.cols)
num_of_cells_needed = int(ceil(width / cw))
if num_of_cells_needed > ss.cols:
w, h = fit_image(width, height, ss.width, height)
ch = int(ss.height / ss.rows)
num_of_rows_needed = int(ceil(height / ch))
y_off = height % ch
cmd['c'], cmd['r'] = ss.cols, num_of_rows_needed
cmd['Y'] = y_off
else:
x_off = width % cw
cmd['X'] = x_off
extra_cells = 0
if align == 'center':
extra_cells = (ss.cols - num_of_cells_needed) // 2
elif align == 'right':
extra_cells = (ss.cols - num_of_cells_needed)
if extra_cells:
sys.stdout.buffer.write(b' ' * extra_cells)
def write_chunked(cmd, data):
if cmd['f'] != 100:
data = zlib.compress(data)
cmd['o'] = 'z'
data = standard_b64encode(data)
while data:
chunk, data = data[:4096], data[4096:]
m = 1 if data else 0
cmd['m'] = m
write_gr_cmd(cmd, chunk)
cmd.clear()
def show(outfile, width, height, fmt, transmit_mode='t', align='center'):
cmd = {'a': 'T', 'f': fmt, 's': width, 'v': height}
set_cursor(cmd, width, height, align)
if detect_support.has_files:
cmd['t'] = transmit_mode
write_gr_cmd(cmd, standard_b64encode(os.path.abspath(outfile).encode(fsenc)))
else:
with open(outfile, 'rb') as f:
data = f.read()
if transmit_mode == 't':
os.unlink(outfile)
if fmt == 100:
cmd['S'] = len(data)
write_chunked(cmd, data)
ImageData = namedtuple('ImageData', 'fmt width height mode')
def run_imagemagick(path, cmd, keep_stdout=True):
try:
p = subprocess.run(cmd, stdout=subprocess.PIPE if keep_stdout else subprocess.DEVNULL, stderr=subprocess.PIPE)
except FileNotFoundError:
raise SystemExit('ImageMagick is required to cat images')
if p.returncode != 0:
raise OpenFailed(path, p.stderr.decode('utf-8'))
return p
def identify(path):
p = run_imagemagick(path, ['identify', '-format', '%m %w %h %A', path])
parts = tuple(filter(None, p.stdout.decode('utf-8').split()))
mode = 'rgb' if parts[3].lower() == 'false' else 'rgba'
return ImageData(parts[0].lower(), int(parts[1]), int(parts[2]), mode)
def convert(path, m, ss):
width, height = m.width, m.height
cmd = ['convert', '-background', 'none', path]
if width > ss.width:
width, height = fit_image(width, height, ss.width, height)
cmd += ['-resize', '{}x{}'.format(width, height)]
with NamedTemporaryFile(prefix='icat-', suffix='.' + m.mode, delete=False) as outfile:
run_imagemagick(path, cmd + [outfile.name])
return outfile.name, width, height
def process(path, args):
m = identify(path)
ss = screen_size()
needs_scaling = m.width > ss.width
if m.fmt == 'png' and not needs_scaling:
outfile = path
transmit_mode = 'f'
fmt = 100
width, height = m.width, m.height
else:
fmt = 24 if m.mode == 'rgb' else 32
transmit_mode = 't'
outfile, width, height = convert(path, m, ss)
show(outfile, width, height, fmt, transmit_mode, align=args.align)
print() # ensure cursor is on a new line
def scan(d):
for dirpath, dirnames, filenames in os.walk(d):
for f in filenames:
mt = mimetypes.guess_type(f)[0]
if mt and mt.startswith('image/'):
yield os.path.join(dirpath, f), mt
def detect_support(wait_for=10):
fd = sys.stdin.fileno()
old = termios.tcgetattr(fd)
oldfl = fcntl.fcntl(fd, fcntl.F_GETFL)
fcntl.fcntl(fd, fcntl.F_SETFL, oldfl | os.O_NONBLOCK)
print('Checking for graphics ({}s max. wait)...'.format(wait_for), end='\r')
sys.stdout.flush()
tty.setraw(fd)
try:
received = b''
start_time = monotonic()
responses = {}
def parse_responses():
for m in re.finditer(b'\033_Gi=([1|2]);(.+?)\033\\\\', received):
iid = m.group(1)
if iid in (b'1', b'2'):
iid = int(iid.decode('ascii'))
if iid not in responses:
responses[iid] = m.group(2) == b'OK'
def read():
nonlocal received
d = sys.stdin.buffer.read()
if not d: # EOF
responses[1] = responses[2] = False
return
received += d
parse_responses()
with NamedTemporaryFile() as f:
f.write(b'abcd'), f.flush()
write_gr_cmd(dict(a='q', s=1, v=1, i=1), standard_b64encode(b'abcd'))
write_gr_cmd(dict(a='q', s=1, v=1, i=2, t='f'), standard_b64encode(f.name.encode(fsenc)))
sel = selectors.DefaultSelector()
sel.register(sys.stdin, selectors.EVENT_READ, read)
while monotonic() - start_time < wait_for and 1 not in responses and 2 not in responses:
for key, mask in sel.select(0.1):
read()
finally:
sys.stdout.buffer.write(b'\033[J'), sys.stdout.flush()
termios.tcsetattr(fd, termios.TCSADRAIN, old)
fcntl.fcntl(fd, fcntl.F_SETFL, oldfl)
detect_support.has_files = bool(responses.get(2))
return responses.get(1, False)
def main(args=sys.argv):
signal.signal(signal.SIGWINCH, lambda: screen_size(refresh=True))
if not sys.stdout.isatty() or not sys.stdin.isatty():
raise SystemExit(
'Must be run in a terminal, stdout is currently not a terminal'
)
if screen_size().width == 0:
raise SystemExit(
'Terminal does not support reporting screen sizes via the TIOCGWINSZ ioctl'
)
msg = (
'A cat like utility to display images in the terminal.'
' You can specify multiple images files and/or directories.'
' Directories are scanned recursively for image files.')
args, items = parse_args(args[1:], options_spec, 'image-file ...', msg, '{} icat'.format(appname))
if not items:
raise SystemExit('You must specify at least one file to cat')
if not detect_support():
raise SystemExit('This terminal emulator does not support the graphics protocol, use a terminal emulator such as kitty that does support it')
errors = []
for item in items:
try:
if os.path.isdir(item):
for x in scan(item):
process(item, args)
else:
process(item, args)
except OpenFailed as e:
errors.append(e)
if not errors:
return
for err in errors:
print(err, file=sys.stderr)
raise SystemExit(1)
if __name__ == '__main__':
main()