i3status script

i3status-maze.py

#!/usr/bin/env python2
# -*- coding: utf8; -*-

import json
import glob
import multiprocessing
import os
import sys
import time

try:
    import alsaaudio
except ImportError:
    alsaaudio = None


DEFAULT_COLOR = '#CCCCCC'
DEFAULT_ICONS = dict(
    battery     = unichr(0x26a1),
    date        = unichr(0x25f7),
    disk        = unichr(0x26c3),
    load        = unichr(0x2605),
    memory      = unichr(0x2631),
    network_up  = unichr(0x2191),
    network_dn  = unichr(0x2193),
    pid_up      = unichr(0x221a),
    pid_dn      = unichr(0x00d7),
    volume      = unichr(0x266a),
)
DEFAULT_HBAR = map(unichr, [
    0x258f, 0x258e, 0x258d, 0x258c,
    0x258b, 0x258a, 0x2589, 0x2588,
])
DEFAULT_VBAR = map(unichr, [
    0x2581, 0x2582, 0x2583, 0x2584,
    0x2585, 0x2586, 0x2587, 0x2588,
])

def hbar(percent, size):
    chars = len(DEFAULT_HBAR)
    positions = chars * size
    full, part = divmod(int(round(positions * (percent / 100.0))), chars)
    output = []
    output.extend([DEFAULT_HBAR[-1]] * full)
    if part:
        output.append(DEFAULT_HBAR[part])
        size -= 1
    output.extend(' ' * (size - full))
    return u''.join(output)

def vbar(percent, size=None):  # size is ignored
    return DEFAULT_VBAR[int(8 * (percent / 100.0))]

def cat(filename):
    return file(filename).read()

def dump(data):
    return json.dumps(data, separators=(',', ':'))

def item(name, text, color=DEFAULT_COLOR, icon=None, **kwargs):
    icon = icon or DEFAULT_ICONS.get(name, None)
    if icon is not None:
        text = (u''.join([icon, ' ', text])).encode('utf-8')
    info = dict(name=name, full_text=text, color=color)
    info.update(kwargs)
    return info


class Metric(object):
    def battery(self, device='BAT0'):
        path = '/sys/class/power_supply/%s' % (device,)
        if not os.path.exists(path):
            return item('battery',
                'N/A',
                color='#FF0000',
                instance=device,
            )

        with open(os.path.join(path, 'uevent')) as handle:
            info = {}
            for line in handle:
                key, value = line.split('=', 1)
                info[key.replace('POWER_SUPPLY_', '').lower()] = value

            perc = 100
            if info['status'].lower() == 'full' or info['capacity'] == '100':
                state = 'full'
                color = DEFAULT_COLOR

            else:
                perc = min(100, int(info['capacity']))
                if perc < 10:
                    color = '#FF0000'
                elif perc < 20:
                    color = '#FFFF00'
                else:
                    color = DEFAULT_COLOR
                state = ' %02d%%' % (perc,)

            return item('battery',
                ' '.join([hbar(perc, 5), state]),
                color=color,
            )

    def date(self, format='%Y-%m-%d %H:%M:%S'):
        return item('date',
            time.strftime(format),
            icon=unichr(0x25f7),
        )

    def disk(self, mount='/'):
        info = os.statvfs(mount)
        size = (info.f_blocks * info.f_bsize) / 1024 ** 2
        free = (info.f_bfree * info.f_bsize) / 1024 ** 2
        perc = (float(free) / float(size)) * 100.0

        if perc < 10:
            color = '#FF0000'
        elif perc < 20:
            color = '#FFFF00'
        else:
            color = DEFAULT_COLOR

        return item('disk_info',
            '%d GB %s' % (free / 1024, mount),
            icon=unichr(0x25a4),
            color=color,
            instance=mount,
        )

    def load(self):
        with open('/proc/loadavg') as handle:
            line = handle.read()
            load = float(line.split()[0])
            cpus = multiprocessing.cpu_count()
            if load > (cpus * 2):
                color = '#FF0000'
            elif load > cpus:
                color = '#FFFF00'
            else:
                color = DEFAULT_COLOR
            return item('load', '%.02f' % (load,), color=color)

    def memory(self):
        with open('/proc/meminfo') as handle:
            size = None
            free = None
            for line in handle:
                k, v = line.split(':')
                if k == 'MemTotal':
                    size = int(v.strip().split()[0])
                elif k == 'MemFree':
                    free = int(v.strip().split()[0])

            if size:
                perc = (float(free) / float(size)) * 100.0
                if perc < 10:
                    color = '#FF0000'
                elif perc < 20:
                    color = '#FFFF00'
                else:
                    color = DEFAULT_COLOR
                return item('memory',
                    '%s %d/%d MB RAM' % (vbar(perc, 5), free / 1024, size / 1024),
                    color=color,
                )
            else:
                return item('memory', '?', color='#424242')

    def network(self):
        path = '/sys/class/net'
        for device in glob.glob('%s/*' % path):
            target = os.readlink(os.path.join(path, device))
            if not target.startswith('/'):
                target = os.path.join(path, target)
            if 'virtual/net' in target:
                continue

            iface = os.path.basename(device)
            state = cat(os.path.join(target, 'operstate')).strip()
            color = dict(
                up='#00FF00',
                down='#FF0000',
            ).get(state, '#424242')
            icon = DEFAULT_ICONS.get('network_%s' % state)
            return item('ethernet',
                '%s %s' % (iface, state),
                color=color,
                instance=iface,
                icon=icon,
            )

    def pid(self, filename, name):
        color = '#FF0000'
        icon = 'pid_dn'
        for pidfile in glob.glob(filename):
            try:
                pid = int(cat(pidfile))
            except ValueError:
                continue

            if os.path.exists('/proc/%d' % pid):
                color = '#00FF00'
                icon = 'pid_up'
                break

        return item(name.lower(), name, color=color,
            icon=DEFAULT_ICONS.get(icon))

    def volume_alsa(self, control='Master'):
        if alsaaudio is None:
            return item('volume', 'N/A', color='#FF0000')

        mixer = alsaaudio.Mixer(control)
        perc = mixer.getvolume()[0]
        volume = 'full'
        if mixer.getmute()[0] or perc < 1:
            perc = 0
            volume = 'mute'
        elif perc < 100:
            volume = ' %02d%%' % (perc,)

        return item('volume', ' '.join([hbar(perc, 5), volume]))


enabled = (
    'network',
    'pid|/run/dhcpd-*.pid,DHCP',
    'pid|/run/openvpn-*.pid,VPN',
    'disk|/',
    'load',
    'memory',
    'volume_alsa',
    'battery',
    'date',
)

def run():
    print dump({'version': 1})
    print '['
    first = True
    while True:
        output = []
        metric = Metric()
        for poller in enabled:
            if '|' in poller:
                hook, args = poller.split('|', 1)
                args = args.split(',')
            else:
                hook, args = poller, ()
            output.append(getattr(metric, hook)(*args))

        if not first:
            sys.stdout.write(',')
        else:
            first = False

        sys.stdout.write(dump(output))
        sys.stdout.flush()
        time.sleep(1)

if __name__ == '__main__':
    sys.exit(run())