square-cpu

#!/usr/bin/python3
import time
import curses
import math
import subprocess

import squarify

from helpers_categorizecmd import categorize_cmd

kilo = 1024
mega = kilo*kilo
giga = mega*kilo


def num_cpus(fallback=None):
    # I have a much better version somewhere.
    cprocs = 0
    try:
        f = open('/proc/cpuinfo')
        for line in f:
            if 'processor' in line:
                cprocs += 1
    except:
        raise
    finally:
        f.close()

    if cprocs==0:
        return fallback
    else:
        return cprocs



    
def main():
    import optparse
    p=optparse.OptionParser()
    p.add_option('-i','--interval',  default=1.0,   action="store", dest="interval",
                 help="Interval (sleep time) between measuring/showing.")    
    options,args = p.parse_args()
    interval = float(options.interval)

    
    ncpus = num_cpus()


    screen = curses.initscr()
    try:
        curses.noecho()
        curses.curs_set(0)

        curses.start_color()
        curses.use_default_colors() # 
        num_inited_cols = 0

        initcolnum = min(curses.COLOR_PAIRS, curses.COLORS) # I should check whether this makes any real sesne

        if curses.COLORS in (8,256):
            import syslog
            while num_inited_cols < curses.COLOR_PAIRS:
                for bg in range(8):
                    if bg in (7,6,5,3,2): # brighter colors
                        fg=0 # black
                    else: # darker colors (black, red, blue)
                        fg = 7 # white
                        # CONSIDER: also using bright white, 15

                    if fg!=bg:
                        num_inited_cols += 1
                        curses.init_pair(num_inited_cols, fg, bg)

                    if num_inited_cols > curses.COLOR_PAIRS:
                        break

            syslog.syslog("inited %d color pairs"%num_inited_cols)

        # TODO: figure out whether we can be smarter for 256
        # the issue is that it does a bunch of foreground colors, on few background colors.
        # we may be better off drawing borders instead of background colors?

        #else:
        #    raise ValueError("My color code doesn't make sense")

        while True:       
            if curses.is_term_resized(curses.COLS,curses.LINES):
                y, x = screen.getmaxyx()
                screen.clear()
                curses.resizeterm(y, x)


            p = subprocess.Popen("ps --no-header -eo pid,pcpu,comm",shell=True, stdout=subprocess.PIPE, encoding='utf8')        # rss is non-swapped physical
            out,_ = p.communicate()
            pid_cmd_pcpu = {}
            for line in out.splitlines():
                ll = line.strip().split()
                pid = ll[0]
                pcpu = float(ll[1])
                cmd = ' '.join(ll[2:])
                pid_cmd_pcpu[pid] = (cmd,pcpu)

            pername = {}

            for pid in pid_cmd_pcpu:
                cmd,pcpu = pid_cmd_pcpu[pid]
                name = cmd

                if pcpu==0.0:
                    continue
                name = categorize_cmd(cmd)
                if name not in pername:
                    pername[name]=0
                pername[name]+=pcpu

            # make list, filter, sum everything under 1% into 'small misc', sort
            name_cpu = pername.items()

            pcpu_thresh = 0.005*(100*ncpus) # also to filter out 0
            temp=[]
            restsize=0
            for name, pcpu in name_cpu:
                if pcpu < pcpu_thresh:
                    restsize+=pcpu
                else:
                    temp.append( (name,pcpu) )
            temp.append( ('(sumsmaller)', restsize) )
            name_cpu = temp

            idlepercent = max( 0, # it's slightly creative with the concept of percent
                                   (100.*ncpus) - sum(cpu   for _,cpu in name_cpu) )
            if idlepercent>0:
                name_cpu.append( ('(idle)', idlepercent ) )


            # normalize to 100%, both for rounding error, and because totaling to 600% is weird.
            totalpercent = sum(pcpu   for name,pcpu in name_cpu)

            name_cpu = list( (name,(100.*pcpu)/totalpercent)   for name,pcpu in name_cpu)


            name_cpu.sort(key=lambda x:x[1], reverse=True)

            pcpus = list(pcpu   for name, pcpu in name_cpu)
            # exhaggerate size of the smaller ones somewhat, for readability
            #pcpus = list(math.pow(pcpu,0.7)   for name, pcpu in name_cpu)

            normed = squarify.normalize_sizes(pcpus, curses.COLS,curses.LINES)
            rects = squarify.squarify(normed, 0, 0, curses.COLS,curses.LINES)

            full = curses.newwin(curses.LINES,curses.COLS,0,0)
            full.clear()
            full.refresh()

            for i, rect in enumerate(rects):
                x = rect['x']
                y = rect['y']
                w = rect['dx']
                h = rect['dy']
                #x=int(x)
                #y=int(y)
                #w=int(w)
                #h=int(h)
                #TODO: check whether this math and ceil stuff is correct, I suspect it isn't.
                x=int(math.floor(round(x)))
                y=int(math.floor(round(y)))
                w=int(math.ceil(round(w)))
                h=int(math.ceil(round(h)))
                #helpers_log.error(` x,y,h,w`)
                try:
                    win = curses.newwin(h,w,y,x)
                    name = name_cpu[i][0]
                    try:
                        import binascii
                        checksum = abs(binascii.crc32(name))
                    except:
                        checksum = sum(list(ord(c) for c in name))
                    colnum = 1+ (checksum % num_inited_cols)
                    if name=='(idle)':
                        colnum = 0 # under current color logic this forces it to black
                    #win.clrtobot()
                    #win.border()
                    win.bkgd( ' ', curses.color_pair(colnum) )
                    win.addstr(1,1, name[:max(0,w-1)] )
                    win.addstr(2,1, ('%d%%'%name_cpu[i][1])[:max(0,w-1)]  )
                    win.redrawwin()
                    win.refresh()
                except curses.error:
                    pass # assume it's because we're tring to draw out of screen due to rounding blah I should look at 

            time.sleep(interval)
    finally:
        curses.endwin()
        curses.echo()

if __name__ == '__main__':
    main()