tmuxchooser

#!/usr/bin/env python3
"""
Shows a list of running tmux sessions with metadata, and allows you to connect
to any running session or create a new one.

Much shiny, very tab completion.

Such Copyright 2013-2014 Michael Ensslin. Wow. GPLv3+.
"""

import subprocess
import os
import pwd
import collections
import time
from collections import namedtuple
from types import SimpleNamespace
import shlex

INF = float("+inf")
TMUX = "/usr/bin/tmux"
SHELL = os.environ["SHELL"]


class EnumVal:
    """
    simple named object; designed for use as an enum value.

    e.g.: A = EnumVal('A'); B = EnumVal('B'); f(A)
    """
    # pylint: disable=too-few-public-methods
    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return repr(self.name)


def ttywidth(fileno=1):
    """
    Determines the width of the terminal at the given tty file descriptor.

    Returns infinity for non-terminal fds.
    """
    import fcntl
    import termios
    import struct
    try:
        packed_input = struct.pack('HHHH', 0, 0, 0, 0)
        packed_output = fcntl.ioctl(fileno, termios.TIOCGWINSZ, packed_input)
        _, width, _, _ = struct.unpack('HHHH', packed_output)
        return width
    except OSError:
        return INF


def colorcodesensitive_len(string):
    """
    Determines the printed length of a string, while ignoring control
    characters and SGR (color) escape sequences.

    Might fail horribly for strongs that contain escape sequences other than
    SGR.
    """
    totallen = 0
    inescseq = False

    for char in string:
        if not inescseq:
            if ord(char) >= 0x20:
                totallen += 1
            elif char == '\x1b':
                inescseq = True
        else:
            if 64 <= ord(char) < 127:
                inescseq = False

    return totallen


ALIGN_LEFT = EnumVal("align left")
ALIGN_CENTER = EnumVal("align center")
ALIGN_RIGHT = EnumVal("align right")


def align_string(string, alignment, width):
    """
    Aligns a single string to be exacty 'width' characters wide

    @param alignment
        one of ALIGN_LEFT, ALIGN_CENTER, ALIGN_RIGHT
    """
    length = colorcodesensitive_len(string)
    free = width - length
    if free < 0:
        # we need to cut
        overflow = "..."
        if width <= len(overflow):
            return overflow[:width]
        else:
            return string[:width - len(overflow)] + overflow

    elif alignment == ALIGN_LEFT:
        return string + " " * free

    elif alignment == ALIGN_CENTER:
        return " " * (free // 2) + string + " " * (free - free // 2)

    elif alignment == ALIGN_RIGHT:
        return " " * free + string

    raise Exception("Unknown text alignment: " + str(alignment))


class Column:
    """
    One column of a Table
    """
    def __init__(self, colname, *functors, minw=0, maxw=INF):
        """
        Creates a column.

        @param colname:
            The column's title
        @param functors:
            List of functions to extract column information from row objects.
            Multiple argument types are allowed:
                - single-arg functions that return the column information
                - str (key for getattr)
                - None (-> empty value)
            All functors are attempted, until one returns non-None.
            For allowed return values for the functors, see add_val.
        @param minw:
            Minimum column width. The column will have an automatic size,
            but will always be at least minw chars wide.
        @param maxw:
            Maximum column width. Guess what.
        """
        self.name = colname
        self.width = minw
        self.minw = minw
        self.maxw = maxw
        self.vals = []

        # process functors
        if not functors:
            functors = [colname]

        def normalize_functor(functor):
            """
            converts all allowed values of 'functor' to a callable function
            """
            if functor is None:
                return lambda o: ""
            elif isinstance(functor, str):
                return lambda o: getattr(o, functor)
            elif callable(functor):
                return functor
            else:
                raise TypeError("functor has invalid type: " + repr(functor))

        self.functors = [normalize_functor(functor) for functor in functors]

        # pointer back to owner table
        self.table = None

    def add_val(self, val):
        """
        Adds a value to a column.

        @val:
            The value object, as returned by a functor.

            Allowed types are:
             - str (newlines separate lines)
             - tuple of (string, alignment)
             - list/generator, containing one such tuple per element

            Single lines must be of type str.
        """
        val = self._normalize_val(val)

        # update own width.
        # wlimit is enforced at printing time, when alignment is done.
        self.width = max([self.width] + [len(s[0]) for s in val])

        # append the val to the col
        self.vals.append(val)

    def _normalize_val(self, val):
        """
        Normalizes val to a list of (line, alignment) tuples.
        """
        # convert v to a list, if it contains only a single val
        # that could be a string, or a (string, alignment) tuple.
        if isinstance(val, str):
            val = val.split('\n')
        elif isinstance(val, tuple):
            val = [val]
        else:
            val = list(val)

        if len(val) > self.table.maxrowh:
            # val has too many lines; crop and add '...' at end
            val = val[0:self.table.maxrowh - 1] + [('...', ALIGN_LEFT)]

        # add alignment to list entries
        return [l if isinstance(l, tuple) else (l, ALIGN_LEFT) for l in val]

    def _get_val(self, rowobj):
        """
        Calculates the field value from rowobj using colval.
        """
        for functor in self.functors:
            val = functor(rowobj)
            if val is not None:
                return val

        # none of the functors worked
        raise Exception("Could not generate valid data for field " + self.name)

    def add_obj(self, rowobj):
        """
        Adds the row object to the column.
        """
        self.add_val(self._get_val(rowobj))


class Table:
    """
    Pretty-printer for tables
    """
    def __init__(self, minrowh=1, maxrowh=INF):
        self.minrowh = minrowh
        self.maxrowh = maxrowh

        # stores the row objects for all but the title row
        self.rowobjs = []

        self.cols = []

    def add_col(self, col):
        """
        Adds a table column.
        """
        col.table = self
        self.cols.append(col)

        # store the colname as the first value
        col.add_val((col.name, ALIGN_CENTER))

        # catch up on all already-existing row objects
        for rowobj in self.rowobjs:
            col.add_obj(rowobj)

    def add_row(self, rowobj):
        """
        Adds a row to the table, using rowobj as content source.
        """
        self.rowobjs.append(rowobj)
        for col in self.cols:
            col.add_obj(rowobj)

    def print_row(self, rowid):
        """
        Pretty-prints row #row of the table.
        """
        rowheight = max(len(col.vals[rowid]) for col in self.cols)

        for i in range(rowheight):
            self.print_rowline(rowid, i)

    def print_rowline(self, rowid, lineno):
        """
        Pretty-prints a single line of row #row.
        """
        fields = []

        for col in self.cols:
            try:
                text, alignment = col.vals[rowid][lineno]
            except IndexError:
                # vertical padding
                text, alignment = "", ALIGN_LEFT

            fields.append(align_string(text, alignment, col.width))

        print("\u2502" + "\u2502".join(fields) + '\u2502')

    def print_rowsep(self, left, middle, right):
        """
        Prints a row separator
        """
        print(
            left +
            middle.join("\u2500" * col.width for col in self.cols) +
            right)

    def print(self):
        """
        Pretty-prints the table.
        """
        # calculate column widths
        freespace = ttywidth()
        colspace = 0
        shrinkablespace = 0

        for col in self.cols:
            # limit col width to its own max
            col.width = min(col.width, col.maxw)
            colspace += col.width
            shrinkablespace += col.width - col.minw

        # additional column-shrinking is necessary if the whole table is
        # wider than the tty
        if freespace < 1 + len(self.cols) + colspace:
            needed = colspace + 1 + len(self.cols) - freespace
            if shrinkablespace >= needed:
                # we need to shrink the shrinkable colspace by 'needed' chars
                factor = 1 - (needed / shrinkablespace)

                import math
                for col in self.cols:
                    newwf = col.minw + (col.width - col.minw) * factor
                    neww = math.ceil(newwf)
                    # the lower the fraction, the more easily we could
                    # sacrifice an aditional character in this col
                    col.fraction = newwf - neww
                    needed -= (col.width - neww)
                    col.width = neww

                while needed > 0:
                    victim = None
                    for col in self.cols:
                        if col.width > col.minw:
                            if victim is None:
                                victim = col
                            elif col.fraction < victim.fraction:
                                victim = col

                    if victim is None:
                        # shrinking has failed; we couldn't find a victim.
                        break

                    victim.width -= 1
                    needed -= 1
                    victim.fraction = 1
            else:
                # we can not shrink the table enough. print it anyway
                pass

        self.print_rowsep('\u250c', '\u252c', '\u2510')
        self.print_row(0)
        for i in range(1, len(self.rowobjs) + 1):
            self.print_rowsep('\u251c', '\u253c', '\u2524')
            self.print_row(i)
        self.print_rowsep('\u2514', '\u2534', '\u2518')


ProcessInfo = namedtuple("ProcessInfo", ("ppid", "execname",
                                         "fgpgroup", "pgroup", "cmdline"))


def get_process_info(pid):
    """
    Retrieves info on one POSIX process.
    """
    # read stats from /proc/
    with open("/proc/" + str(pid) + "/stat") as statfile:
        statstr = statfile.read()

    # parsing /proc/stat is horrible.
    leftmiddle, right = statstr.rsplit(')', 1)
    left, middle = leftmiddle.split('(', 1)
    left = left.strip().split(' ')
    middle = [middle]
    right = right.strip().split(' ')
    stats = left + middle + right

    with open("/proc/" + str(pid) + "/cmdline") as cmdlinefile:
        argv = cmdlinefile.read().split('\0')
        if not argv[-1]:
            argv.pop()

    # for field numbers, and why the above hat to be done, see man 5 proc
    return ProcessInfo(
        ppid=int(stats[3]),
        execname=stats[1],
        fgpgroup=int(stats[7]),
        pgroup=int(stats[4]),
        cmdline=" ".join(shlex.quote(arg) for arg in argv))


def list_processes():
    """
    Collects information on every running process

    @returns a dictionary of {pid: ProcessInfo}
    """
    procs = collections.OrderedDict()

    # collect information for all processes
    for pidstring in os.listdir("/proc"):
        try:
            pid = int(pidstring)
            procs[pid] = get_process_info(pid)
        except ValueError:
            pass
        except FileNotFoundError:
            pass

    return procs


def popen(args):
    """
    Runs the given process and returns its stdout.

    @param args:
        process argv
    """
    proc = subprocess.Popen(args,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    if proc.wait() != 0:
        raise RuntimeError(proc.stderr.read().decode('utf-8'))
    return proc.stdout.read().decode('utf-8', errors='replace')


def tmux_get(command, fields):
    """
    Queries tmux with the given command, returning a named tuple of 'fields'.

    @param command:
        str consisting of whitespace-joined argv[1:].
    """
    fstring = ""
    for field in fields:
        fstring += "#{" + field + "}\n"

    stdout = popen([TMUX] + command.split(' ') + ["-F", fstring])
    for item in stdout.split("\n\n"):
        if item == "":
            return

        yield namedtuple("FieldsType", fields)(*item.split('\n'))


def tmux_get_sessions():
    """
    Queries tmux for the currently-running sessions.

    @returns:
        A list of SessionInfo items.
    """

    procs = list_processes()

    sessions = {}
    # retrieve basic session list
    for session in tmux_get("list-sessions", ["session_name",
                                              "session_created",
                                              "session_width",
                                              "session_height"]):

        session_info = SimpleNamespace()

        session_info.name = session.session_name
        session_info.created = int(session.session_created)
        session_info.w = int(session.session_width),
        session_info.h = int(session.session_height),
        session_info.dims = (session.session_width + "x" +
                             session.session_height)
        session_info.fgprocs = []
        session_info.users = {}
        session_info.usr_blame_w = session_info.usr_blame_h = None
        session_info.usr_most_recent = None

        sessions[session_info.name] = session_info

    # retrieve info on foreground process
    for pane in tmux_get("list-panes -a", ["session_name", "pane_pid"]):
        session_info = sessions[pane.session_name]
        pid = int(pane.pane_pid)
        try:
            fgpgroup = procs[pid].fgpgroup
            for pid in procs:
                if procs[pid].pgroup == fgpgroup:
                    session_info.fgprocs.append(procs[pid])
        except KeyError:
            # advertised process doesn't exist
            pass

    # retrieve info on connected clients
    for client in tmux_get("list-clients", ["session_name",
                                            "client_activity",
                                            "client_tty",
                                            "client_width",
                                            "client_height"]):

        session_info = sessions[client.session_name]

        # the user name is the owner of the pts
        uname = pwd.getpwuid(os.stat(client.client_tty).st_uid).pw_name
        last_activity = int(client.client_activity)
        client_width = int(client.client_width)
        client_height = int(client.client_height)
        if uname not in session_info.users:
            usr = SimpleNamespace()

            usr.w = client_width,
            usr.h = client_height,
            usr.last_activity = last_activity

            session_info.users[uname] = usr
        else:
            usr = session_info.users[uname]
            usr.w = min(client_width, usr.w)
            usr.h = min(client_height, usr.h)
            usr.last_activity = max(last_activity, usr.last_activity)

        # check which user is to blame for the terminal size constraints,
        # and who was most recently active
        if usr.w == session_info.w:
            session_info.usr_blame_w = usr
        if usr.h == session_info.h:
            session_info.usr_blame_h = usr
        if session_info.usr_most_recent is None:
            session_info.usr_most_recent = usr
        else:
            if usr.last_activity > session_info.usr_most_recent.last_activity:
                session_info.usr_most_recent = usr

    return list(sessions.values())


def tmux_attach_session(name):
    """
    Attaches to the running session with the given name.
    """
    return subprocess.call([TMUX, "attach-session", "-t", name]) == 0


def tmux_create_session(name):
    """
    Creates a new session with the given name.
    """
    return subprocess.call([TMUX, "new-session", "-s", name]) == 0


class Completer:
    """
    Readline completer for available tmux sessions.
    """
    # pylint: disable=too-few-public-methods

    def __init__(self, sessions):
        self.sessions = [s.name for s in sessions]

        # caches the list of completion options.
        self.cache = []

    def __call__(self, text, status):
        if status == 0:
            self.cache = [s for s in self.sessions if s.startswith(text)]

        try:
            return self.cache[status]
        except IndexError:
            return None


def date_to_str(date):
    """
    Very-human-readable date-to-string function.
    """
    deltatime = round(time.time()) - date
    localtime = time.localtime(date)

    if deltatime < 10:
        return "just now"
    elif deltatime < 60:
        return "last min"
    elif deltatime < (3600 * 12):
        return time.strftime("%H:%M", localtime)
    else:
        return time.strftime("%Y-%m-%d %H:%M", localtime)


def print_shiny_session_table(sessions):
    """
    Prints a shiny table of the given sessions.
    """
    table = Table(minrowh=1, maxrowh=10)

    table.add_col(Column(
        "name",
        minw=10))

    table.add_col(Column(
        "dims",
        lambda session: (session.dims, ALIGN_RIGHT),
        minw=6))

    table.add_col(Column(
        "created",
        lambda session: (date_to_str(session.created), ALIGN_RIGHT),
        minw=16))

    table.add_col(Column(
        "last active",
        lambda session: "" if session.usr_most_recent is None else None,
        lambda session: (date_to_str(session.usr_most_recent.last_activity),
                         ALIGN_RIGHT),
        minw=16))

    table.add_col(Column(
        "users",
        lambda session: (
            (x, ALIGN_RIGHT) for x in
            sorted(session.users,
                   key=lambda x: -session.users[x].last_activity)),
        minw=8))

    table.add_col(Column(
        "fg processes",
        lambda session: [p.cmdline for p in session.fgprocs]))

    for session in sessions:
        table.add_row(session)

    table.print()


def loop():
    """
    One iteration of the main loop.

    Gathers information, queries the user for input, and responds accordingly.
    """
    # pylint: disable=broad-except
    try:
        sessions = tmux_get_sessions()
    except RuntimeError as ex:
        sessions = []
        print(ex)
    except Exception:
        import traceback
        print("Error while retrieving session list:")
        traceback.print_exc()

    print("Running TMUX sessions:\n")
    print_shiny_session_table(sessions)

    readline.set_completer(Completer(sessions))

    print("\n"
          "Type the name of an existing or new session, "
          "or press ENTER to spawn a shell.")

    line = input("> ")
    if not line:
        print("Spawning shell")
        subprocess.call([SHELL])
        return

    if line in (s.name for s in sessions):
        if not tmux_attach_session(line):
            print("Could not attach to session " + line)
    else:
        if not tmux_create_session(line):
            print("Could not create session " + line)


if __name__ == '__main__':
    import readline
    readline.parse_and_bind("tab: complete")

    os.chdir(os.environ['HOME'])

    try:
        while True:
            loop()
    except EOFError:
        pass
    except KeyboardInterrupt:
        pass

    print("")