cputemp

#!/usr/bin/env python3
# cputemp -- visually monitor CPU temperature
import argparse
from collections import defaultdict
from glob import glob
import math
import os
import signal
import sys
import time

def fround(val):
    # Round 0.5 up (Python rounds 0.5 down)
    return round(val + 0.1)

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

def set_title(text):
    print("\033]0;%s\033\\" % text, end="", file=sys.stderr, flush=True)

def set_wname(text):
    if os.environ["TERM"].startswith("tmux"):
        print("\033k%s\033\\" % text, end="", file=sys.stderr, flush=True)

def get_loadavg():
    with open("/proc/loadavg", "r") as fh:
        a, b, c, *_ = fh.read().split()
    return float(a), float(b), float(c)

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 rsuffix(string, old, new):
    if string.endswith(old):
        return string[:-len(old)] + str(new)
    return string

class Sensor:
    def __init__(self, dev, subdev):
        self.dev = dev
        self.subdev = subdev

    @property
    def devname(self):
        with open(f"{self.dev}/name", "r") as fh:
            return fh.read().strip()

    @property
    def label(self):
        path = rsuffix(self.subdev, "_input", "_label")
        with open(path, "r") as fh:
            return fh.read().strip()

    def read(self):
        with open(self.subdev, "r") as fh:
            temp = int(fh.read())
            return temp / 1000

    @classmethod
    def find_hwmon_by_name(cls, name):
        for dev in glob("/sys/class/hwmon/hwmon*"):
            if os.path.exists("%s/name" % dev):
                with open("%s/name" % dev, "r") as fh:
                    dev_name = fh.read().strip()
                if dev_name == "coretemp":
                    return dev
        raise RuntimeError("no %r hwmon device found" % name)

    @classmethod
    def find_sensor(cls):
        dev = cls.find_hwmon_by_name("coretemp")

        subdev = None
        subdev_pkg = None
        subdev_cr0 = None
        for path_label in glob("%s/temp*_label" % dev):
            path_input = rsuffix(path_label, "_label", "_input")
            with open(path_label, "r") as fh:
                label = fh.read().strip()
            if label == "Package id 0":
                subdev_pkg = path_input
            elif label == "Core 0":
                subdev_cr0 = path_input
        if subdev_pkg:
            subdev = subdev_pkg
        elif subdev_cr0:
            subdev = subdev_cr0
        else:
            raise RuntimeError("no 'coretemp/Core 0' hwmon sensor found")

        with open("%s/name" % dev, "r") as fh:
            dev_name = fh.read().strip()
        with open(rsuffix(subdev, "_input", "_label"), "r") as fh:
            subdev_name = fh.read().strip()

        return Sensor(dev, subdev)

def print_histogram(data, *, width=5, unit="", barwidth=50):
    def buckets(data, width):
        res = defaultdict(int)
        for k, v in data.items():
            lo = k - (k % width)
            hi = lo + width
            res[lo, hi] += v
        return res

    total = sum(data.values())
    grouped = buckets(data, width)
    for (lo, hi), v in sorted(grouped.items()):
        percentage = 100 * v / total
        width = barwidth * percentage / 100
        label = f"[{lo:-2d}, {hi:-2d})"
        bar = "=" * round(width)
        suffix = f"({v}{unit}, {percentage:.1f}%)"
        print(f"{label:9} | {bar}  {suffix}")

    bar = "-" * barwidth
    print(f"{'':9} +-{bar}")

class MeterBar:
    def __init__(self, interval=1.0):
        # Last values for redraw/overdraw
        self.last_ts = 0
        self.last_temp = 0
        self.last_load = 0

        # Statistics
        self.interval = interval
        self.max_temp = 0
        self.time_temps = defaultdict(int)

    def _ansicolor_for_temp(self, temp):
        low, high = 25, 75
        stops = [
            (0, 0, 5), # blue
            (0, 5, 0), # green
            (5, 5, 0), # yellow
            (5, 0, 0), # red
        ]

        t = clamp(temp, low, high)
        t = normalize(t, low, high)
        r, g, b = [round(x) for x in linear_gradient(stops, t, gamma=2.2)]

        color = 16 + r*36 + g*6 + b
        return f"38;5;{color}"

    def _bar(self, width, chars, color):
        if width <= 0:
            return ""

        bar = chars[0] * (width-1) + chars[-1]
        if color:
            bar = f"\033[{color}m{bar}\033[m"
        return bar

    def _draw_bar(self, is_current, ts, temp, load):
        width = int(temp) - 20
        width = clamp(width, 1, 60)
        chars = "█" if is_current else "-║"
        color = self._ansicolor_for_temp(temp)

        put(time.strftime("%T ", time.localtime(ts)))
        put(self._bar(width, chars, color))
        put(f" {temp:3d}°C")
        put(f"  \033[2m{load:.2f}\033[m")
        print(flush=True)

    def overdraw_bar(self):
        put("\033[A") # One line up
        put("\r\033[K") # Reset position after signal
        self._draw_bar(False, self.last_ts, self.last_temp, self.last_load)

    def update(self, ts, temp, load):
        self._draw_bar(True, ts, temp, load)

        # Store values for overdraw
        self.last_ts = ts
        self.last_temp = temp
        self.last_load = load

        # Update stats
        self.max_temp = max(self.max_temp, temp)
        self.time_temps[fround(temp)] += args.interval

    def print_stats(self, from_signal=False):
        if from_signal:
            self.overdraw_bar()

        print()
        print("Statistics:")
        print_histogram(self.time_temps, unit="s")
        for _min in [60, 70, 80, 90]:
            temp_secs = sum([_time
                             for (_temp, _time) in self.time_temps.items()
                             if _temp >= _min])
            if temp_secs:
                print(f"Spent {temp_secs}s at or above {_min}°C")

        max_temp = self.max_temp
        max_temp_secs = self.time_temps[max_temp]
        print(f"Maximum temperature: {max_temp}°C (for {max_temp_secs}s)")

        if from_signal:
            print()
            self._draw_bar(True, self.last_ts, self.last_temp, self.last_load)

parser = argparse.ArgumentParser()
parser.add_argument("-1", "--once", action="store_true",
                    help="output current value and exit")
parser.add_argument("-r", "--raw", action="store_true",
                    help="output only temperature value")
parser.add_argument("-s", "--spark", action="store_true",
                    help="output through 'spark'")
parser.add_argument("-l", "--load", action="store_true",
                    help="monitor load average")
parser.add_argument("-W", "--wname", action="store_true",
                    help="update tmux window name")
parser.add_argument("-n", "--interval", metavar="SECS", type=int, default=1,
                    help="specify update interval")
args = parser.parse_args()

hostname = os.uname().nodename
sensor = Sensor.find_sensor()

if args.spark:
    if args.once:
        exit("cputemp: Options --once and --spark are incompatible")
    args.raw = True
    sys.stdout = os.popen("spark -l 20 -h 70", "w")

if args.raw:
    try:
        while True:
            temp = sensor.read()
            load = get_loadavg()[0]

            print(f"{load:.2f}" if args.load else f"{temp:.0f}",
                  flush=True)

            set_title(f"{hostname}: {temp:.0f}°C (load {load:.2f})")
            if args.wname:
                set_wname(f"{load:.2f}" if args.load else f"{temp:.0f}°C")

            if args.once:
                break
            time.sleep(args.interval)
    except KeyboardInterrupt:
        exit(0)
else:
    if args.load:
        print("cputemp: Option --load is ignored when not in raw mode",
              file=sys.stderr)
    print(f"Using sensor {sensor.subdev} ({sensor.devname}:{sensor.label})",
          file=sys.stderr)
    m = MeterBar(interval=args.interval)
    try:
        signal.signal(signal.SIGQUIT, lambda s, f: m.print_stats(True))
        while True:
            if m.last_temp:
                m.overdraw_bar()
            ts = time.time()
            temp = sensor.read()
            temp = fround(temp)
            load = get_loadavg()[0]
            m.update(ts, temp, load)
            set_title(f"{hostname}: {temp:.0f}°C (load {load:.2f})")
            set_wname(f"{temp:.0f}°C")
            if args.once:
                break
            time.sleep(args.interval)
    except RuntimeError as e:
        exit(f"error: {e}")
    except KeyboardInterrupt:
        m.print_stats()
        exit(0)