system-status/system-status.py

import psutil
import sys

from gpiozero import Button
from signal import signal, SIGTERM
from subprocess import check_call
from time import sleep
from unicornhatmini import UnicornHATMini, BUTTON_A

POWER_BUTTON_HOLD_TIME_IN_SECONDS = 1
DISPLAY_UPDATE_IN_SECONDS         = 2

BLACK  = [0,  0,  0 ]
RED    = [64, 0,  0 ]
BLUE   = [0,  16, 32]
ORANGE = [64, 16, 0 ]
PINK   = [64, 0,  16]

TEMP_COLOR    = RED
CPU_COLOR     = BLUE
NETWORK_COLOR = PINK

MIN_TEMP = 40
MAX_TEMP = 85


class SystemStatus:

    def __init__(self, hat):
        self.hat = hat
        self.hat.set_brightness(0.1)
        self.hat.set_rotation(270)
        self.width, self.height = self.hat.get_shape()

        self.tempStartColumn    = 0
        self.tempEndColumn      = int(self.width / 2)
        self.cpuStartColumn     = int(self.width / 2)
        self.cpuEndColumn       = self.width - 1
        self.networkStartColumn = self.width - 1
        self.networkEndColumn   = self.width

        self.start()

    def terminate(self):
        self.stop()

        for _ in range(3):
            sleep(0.25)
            self.hat.set_all(*RED)
            self.hat.show()
            sleep(0.25)
            self.clear()

    def start(self):
        self.isRunning = True
        self.clear()

    def stop(self):
        self.isRunning = False
        self.clear()

    def clear(self):
        self.hat.set_all(*BLACK)
        self.hat.show()

    def display(self):
        if self.isRunning:
            self.displayTemp()
            self.displayCpu()
            self.displayNetwork()
            self.hat.show()

    def displayTemp(self):
        self.displayField(
            start=self.tempStartColumn,
            end=self.tempEndColumn,
            value=self.getTempValue(),
            color=TEMP_COLOR
        )

    def displayCpu(self):
        self.displayField(
            start=self.cpuStartColumn,
            end=self.cpuEndColumn,
            value=self.getCpuValue(),
            color=CPU_COLOR
        )

    def displayNetwork(self):
        self.displayField(
            start=self.networkStartColumn,
            end=self.networkEndColumn,
            value=self.getNetworkValue(),
            color=NETWORK_COLOR
        )

    def displayField(self, start, end, value, color):
        for x in range(start, end):
            for y in range(value):
                self.hat.set_pixel(x, y, *color)
            for y in range(value, self.height):
                self.hat.set_pixel(x, y, *BLACK)

    def getTempValue(self):
        temp = psutil.sensors_temperatures()['cpu_thermal'][0].current
        adjusted = max(0, temp - MIN_TEMP)
        percent = min(1, adjusted / (MAX_TEMP - MIN_TEMP))

        return max(1, round(self.height * percent))

    def getCpuValue(self):
        percent = psutil.cpu_percent() / 100

        return round(self.height * percent)

    def getNetworkValue(self):
        connections = psutil.net_connections(kind='inet')
        remoteConnections = [c for c in connections if c.raddr]

        return min(self.height, len(remoteConnections))


def main():
    systemStatus = SystemStatus(UnicornHATMini())
    powerButton = Button(BUTTON_A, hold_time=POWER_BUTTON_HOLD_TIME_IN_SECONDS)
    signal(SIGTERM, cleanExit(systemStatus, powerButton))

    powerButton.when_held = poweroff(systemStatus)

    while True:
        systemStatus.display()
        sleep(DISPLAY_UPDATE_IN_SECONDS)


def cleanExit(systemStatus, button):
    def _cleanExit_(signum, frame):
        button.close()
        systemStatus.stop()
        sys.exit(0)

    return _cleanExit_


def poweroff(systemStatus):
    def _poweroff_():
        systemStatus.terminate()
        check_call(['sudo', 'poweroff'])

    return _poweroff_


if __name__ == '__main__':
    main()