osc_yank.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import base64
import os
import subprocess
import sys


PY2 = sys.version_info.major < 3  # is needed for correct mypy checking
MAX_OSC_YANK_LEN = 74994

if PY2:
    str = unicode  # pylint: disable=redefined-builtin
    stdin_bytes = sys.stdin
    stdout_bytes = sys.stdout
else:
    stdin_bytes = sys.stdin.buffer
    stdout_bytes = sys.stdout.buffer


class NoTtyError(Exception):
    pass


def get_active_tmux_tty():
    # type: () -> str
    try:
        tmux_active_tty = [
            tty
            for is_active, tty in (
                line.split()
                for line in subprocess.check_output(["tmux", "list-panes", "-F", "#{pane_active} #{pane_tty}"])
                .strip()
                .split(b"\n")
            )
            if is_active
        ][0]
    except (subprocess.CalledProcessError, IndexError):
        raise NoTtyError
    return tmux_active_tty


def osc_yank(data):
    # type: (bytes) -> None
    if len(data) > MAX_OSC_YANK_LEN:
        print("Input is {} bytes too long".format(len(data) - MAX_OSC_YANK_LEN), file=sys.stderr)
        data = data[:MAX_OSC_YANK_LEN]

    osc_sequence = b"\033]52;c;" + base64.b64encode(data) + b"\a"  # type: bytes
    if "TMUX" not in os.environ:
        stdout_bytes.write(osc_sequence)
    else:
        # The tmux `run-shell` command does not allow to write escape codes to stdout
        # -> Get the used tty and write to the tty directly
        tmux_tty = get_active_tmux_tty()
        with open(tmux_tty, "wb") as tty:
            tty.write(osc_sequence)


def main():
    # type: () -> None
    data = stdin_bytes.read()
    try:
        osc_yank(data)
    except NoTtyError:
        sys.exit(2)


if __name__ == "__main__":
    main()