spark

#!/usr/bin/env python3
# spark -- output a "spark line" graph
#
#     $ seq 0 7 | spark -l 0 -h 7
#     ▁▂▃▄▅▆▇█ 7.0
#
# Reads a series of numbers from stdin and outputs them as a "spark line",
# similar to how `prettyping` works. Input is assumed to be streamed, so the
# low/high bounds must be specified explicitly:
#
#     while sleep 1; do get_temp; done | spark --low 30 --high 50
#
# Note to self: It appears impossible to support showing summary on Ctrl-\
# (like in `ping` or `cputemp`), as the SIGQUIT will be sent to the entire
# pipeline and will kill the data source process.

import argparse
import datetime
import math
import os
import signal
import sys
import termios

def clamp(val, low, high):
    return min(max(val, low), high)

def normalize(val, low, high):
    return (val - low) / (high - low)

def mix(a, b, t):
    return a*(1-t) + b*t

def vec_mix(a, b, t):
    assert len(a) == len(b)
    return [mix(a[i], b[i], t) for i in range(len(a))]

def vec_gamma_scale(a, gamma):
    return [x**gamma for x in a]

def linear_gradient(stops, value, gamma=1.0):
    n = 1 / (len(stops) - 1)
    ratio = value / n
    i = math.floor(ratio)
    if i == len(stops) - 1:
        return stops[len(stops) - 1]
    frac = ratio % 1
    stops = [vec_gamma_scale(x, gamma) for x in stops]
    result = vec_mix(stops[i], stops[i+1], frac)
    return vec_gamma_scale(result, 1/gamma)

def put(s):
    print(s, end="", flush=True)

class no_echo:
    def __enter__(self):
        self.old_attrs = termios.tcgetattr(sys.stdout)
        attrs = self.old_attrs[:]
        attrs[3] &= ~termios.ECHO
        termios.tcsetattr(sys.stdout, termios.TCSAFLUSH, attrs)
        return self

    def __exit__(self, *_):
        termios.tcsetattr(sys.stdout, termios.TCSAFLUSH, self.old_attrs)

class Sparkline:
    TICKS = "▁▂▃▄▅▆▇█"

    SIX_COLORS = [
        (0, 0, 5), # blue
        (0, 5, 0), # green
        (5, 5, 0), # yellow
        (5, 0, 0), # red
    ]

    RGB_COLORS = [
        (0x00, 0x00, 0xFF), # blue
        (0x00, 0xFF, 0x00), # green
        (0xFF, 0xFF, 0x00), # yellow
        (0xFF, 0x00, 0x00), # red
    ]

    def __init__(self, low, high, *, sep=False, invert=False):
        self.low = low
        self.high = high
        self.sep = sep

        # Horizontal cursor position
        self.w = 0
        self.update_ttysize()

        self.last_tail = None
        self.last_date = None

        if os.environ.get("SPARK_NO_RGB"):
            self.truecolor = False
            self.colors = self.SIX_COLORS
        else:
            self.truecolor = True
            self.colors = self.RGB_COLORS

        if invert:
            self.colors = self.colors[1:] # remove blue
            self.colors = self.colors[::-1]

    def update_ttysize(self):
        if x := os.environ.get("COLUMNS"):
            self.ncols = int(x)
        else:
            self.ncols = os.get_terminal_size(sys.stderr.fileno()).columns

    def _print_tick(self, val):
        val = clamp(val, self.low, self.high)

        t = normalize(val, self.low, self.high)
        t = (val - self.low) / (self.high - self.low)
        ti = math.floor(t * (len(self.TICKS)-1))
        r, g, b = [round(x) for x in linear_gradient(self.colors, t, gamma=1.8)]

        if self.truecolor:
            put(f"\033[38;2;{r};{g};{b}m")
        else:
            put(f"\033[38;5;{16 + r*36 + g*6 + b}m")
        put(self.TICKS[ti])
        put("\033[m")
        self.w += 1

    def _maybe_print_date(self):
        if not self.sep:
            return

        now = datetime.datetime.now()
        was = (self.last_date or now)

        if (was.date(), was.time().hour) != (now.date(), now.time().hour):
            out = now.strftime("%Y-%m-%d %H:%M")
            out = f"-- {out} ".ljust(self.ncols, "-")
            put(f"\033[38;5;244m{out}\033[m\n")

        self.last_date = now

    def update(self, val):
        tail = f" {val:4.1f}"

        put("\033[K")
        if self.w + len(tail) + 1 >= self.ncols:
            put(f"\033[38;5;244m{self.last_tail}\033[m")
            put("\n")
            self.w = 0
            self._maybe_print_date()

        self._print_tick(val)

        put(f"\033[1m{tail}\033[m")
        put("\b" * len(tail))
        self.last_tail = tail

parser = argparse.ArgumentParser(add_help=False)
parser.description = "Draws a sparkline graph from numeric inputs."
parser.add_argument("--help", action="help",
                    help="show this help message and exit")
parser.add_argument("-l", "--low", type=float, default=0,
                    help="minimum value to clamp the sparkline to")
parser.add_argument("-h", "--high", type=float, default=100,
                    help="maximum value to clamp the sparkline to")
parser.add_argument("-i", "--invert", action="store_true",
                    help="invert color relationship")
parser.add_argument("-n", "--no-sep", action="store_true",
                    help="disable hourly separators")
args = parser.parse_args()

if args.low >= args.high:
    exit("error: Maximum must be higher than minimum")

s = Sparkline(low=args.low,
              high=args.high,
              sep=(not args.no_sep),
              invert=args.invert)

if hasattr(signal, "SIGWINCH"):
    signal.signal(signal.SIGWINCH,
                  lambda signum, stack: s.update_ttysize())

try:
    with no_echo():
        for line in sys.stdin:
            for val in line.split():
                val = float(val)
                s.update(val)
        put("\n")
except KeyboardInterrupt:
    exit()