refactor scripts using curses

3 files changed, 117 insertions, 65 deletions

diff --git a/scripts/iotop.py b/scripts/iotop.py
index c3afd071..04683673 100755
--- a/scripts/iotop.py
+++ b/scripts/iotop.py
@@ -14,7 +14,7 @@ It doesn't work on Windows as curses module is required.
 
 Example output:
 
-$ python scripts/iotop.py
+$ python3 scripts/iotop.py
 Total DISK READ: 0.00 B/s | Total DISK WRITE: 472.00 K/s
 PID   USER      DISK READ  DISK WRITE  COMMAND
 13155 giampao    0.00 B/s  428.00 K/s  /usr/bin/google-chrome-beta
@@ -30,7 +30,6 @@ PID   USER      DISK READ  DISK WRITE  COMMAND
 Author: Giampaolo Rodola' <g.rodola@gmail.com>
 """
 
-import atexit
 import time
 import sys
 try:
@@ -42,20 +41,11 @@ import psutil
 from psutil._common import bytes2human
 
 
-def tear_down():
-    win.keypad(0)
-    curses.nocbreak()
-    curses.echo()
-    curses.endwin()
-
-
 win = curses.initscr()
-atexit.register(tear_down)
-curses.endwin()
 lineno = 0
 
 
-def print_line(line, highlight=False):
+def printl(line, highlight=False):
     """A thin wrapper around curses's addstr()."""
     global lineno
     try:
@@ -129,10 +119,10 @@ def refresh_window(procs, disks_read, disks_write):
 
     disks_tot = "Total DISK READ: %s | Total DISK WRITE: %s" \
                 % (bytes2human(disks_read), bytes2human(disks_write))
-    print_line(disks_tot)
+    printl(disks_tot)
 
     header = templ % ("PID", "USER", "DISK READ", "DISK WRITE", "COMMAND")
-    print_line(header, highlight=True)
+    printl(header, highlight=True)
 
     for p in procs:
         line = templ % (
@@ -142,21 +132,45 @@ def refresh_window(procs, disks_read, disks_write):
             bytes2human(p._write_per_sec),
             p._cmdline)
         try:
-            print_line(line)
+            printl(line)
         except curses.error:
             break
     win.refresh()
 
 
+def setup():
+    curses.start_color()
+    curses.use_default_colors()
+    for i in range(0, curses.COLORS):
+        curses.init_pair(i + 1, i, -1)
+    curses.endwin()
+    win.nodelay(1)
+
+
+def tear_down():
+    win.keypad(0)
+    curses.nocbreak()
+    curses.echo()
+    curses.endwin()
+
+
 def main():
+    global lineno
+    setup()
     try:
         interval = 0
         while True:
+            if win.getch() == ord('q'):
+                break
             args = poll(interval)
             refresh_window(*args)
-            interval = 1
+            lineno = 0
+            interval = 0.5
+            time.sleep(interval)
     except (KeyboardInterrupt, SystemExit):
         pass
+    finally:
+        tear_down()
 
 
 if __name__ == '__main__':
diff --git a/scripts/nettop.py b/scripts/nettop.py
index ce647c9d..8cc19fda 100755
--- a/scripts/nettop.py
+++ b/scripts/nettop.py
@@ -11,7 +11,7 @@ Shows real-time network statistics.
 
 Author: Giampaolo Rodola' <g.rodola@gmail.com>
 
-$ python scripts/nettop.py
+$ python3 scripts/nettop.py
 -----------------------------------------------------------
 total bytes:           sent: 1.49 G       received: 4.82 G
 total packets:         sent: 7338724      received: 8082712
@@ -31,7 +31,6 @@ pkts-sent                     0               0
 pkts-recv               1214470               0
 """
 
-import atexit
 import time
 import sys
 try:
@@ -43,20 +42,11 @@ import psutil
 from psutil._common import bytes2human
 
 
-def tear_down():
-    win.keypad(0)
-    curses.nocbreak()
-    curses.echo()
-    curses.endwin()
-
-
-win = curses.initscr()
-atexit.register(tear_down)
-curses.endwin()
 lineno = 0
+win = curses.initscr()
 
 
-def print_line(line, highlight=False):
+def printl(line, highlight=False):
     """A thin wrapper around curses's addstr()."""
     global lineno
     try:
@@ -89,59 +79,80 @@ def refresh_window(tot_before, tot_after, pnic_before, pnic_after):
     global lineno
 
     # totals
-    print_line("total bytes:           sent: %-10s   received: %s" % (
+    printl("total bytes:           sent: %-10s   received: %s" % (
         bytes2human(tot_after.bytes_sent),
         bytes2human(tot_after.bytes_recv))
     )
-    print_line("total packets:         sent: %-10s   received: %s" % (
+    printl("total packets:         sent: %-10s   received: %s" % (
         tot_after.packets_sent, tot_after.packets_recv))
 
     # per-network interface details: let's sort network interfaces so
     # that the ones which generated more traffic are shown first
-    print_line("")
+    printl("")
     nic_names = list(pnic_after.keys())
     nic_names.sort(key=lambda x: sum(pnic_after[x]), reverse=True)
     for name in nic_names:
         stats_before = pnic_before[name]
         stats_after = pnic_after[name]
         templ = "%-15s %15s %15s"
-        print_line(templ % (name, "TOTAL", "PER-SEC"), highlight=True)
-        print_line(templ % (
+        printl(templ % (name, "TOTAL", "PER-SEC"), highlight=True)
+        printl(templ % (
             "bytes-sent",
             bytes2human(stats_after.bytes_sent),
             bytes2human(
                 stats_after.bytes_sent - stats_before.bytes_sent) + '/s',
         ))
-        print_line(templ % (
+        printl(templ % (
             "bytes-recv",
             bytes2human(stats_after.bytes_recv),
             bytes2human(
                 stats_after.bytes_recv - stats_before.bytes_recv) + '/s',
         ))
-        print_line(templ % (
+        printl(templ % (
             "pkts-sent",
             stats_after.packets_sent,
             stats_after.packets_sent - stats_before.packets_sent,
         ))
-        print_line(templ % (
+        printl(templ % (
             "pkts-recv",
             stats_after.packets_recv,
             stats_after.packets_recv - stats_before.packets_recv,
         ))
-        print_line("")
+        printl("")
     win.refresh()
     lineno = 0
 
 
+def setup():
+    curses.start_color()
+    curses.use_default_colors()
+    for i in range(0, curses.COLORS):
+        curses.init_pair(i + 1, i, -1)
+    curses.endwin()
+    win.nodelay(1)
+
+
+def tear_down():
+    win.keypad(0)
+    curses.nocbreak()
+    curses.echo()
+    curses.endwin()
+
+
 def main():
+    setup()
     try:
         interval = 0
         while True:
+            if win.getch() == ord('q'):
+                break
             args = poll(interval)
             refresh_window(*args)
-            interval = 1
+            interval = 0.5
     except (KeyboardInterrupt, SystemExit):
         pass
+    finally:
+        tear_down()
 
 
 if __name__ == '__main__':
diff --git a/scripts/top.py b/scripts/top.py
index 0b17471d..3c297d9a 100755
--- a/scripts/top.py
+++ b/scripts/top.py
@@ -9,7 +9,7 @@ A clone of top / htop.
 
 Author: Giampaolo Rodola' <g.rodola@gmail.com>
 
-$ python scripts/top.py
+$ python3 scripts/top.py
  CPU0  [||||                                    ]  10.9%
  CPU1  [|||||                                   ]  13.1%
  CPU2  [|||||                                   ]  12.8%
@@ -33,7 +33,6 @@ PID    USER       NI   VIRT    RES  CPU%  MEM%     TIME+  NAME
 ...
 """
 
-import atexit
 import datetime
 import sys
 import time
@@ -46,30 +45,28 @@ import psutil
 from psutil._common import bytes2human
 
 
-# --- curses stuff
-
-def tear_down():
-    win.keypad(0)
-    curses.nocbreak()
-    curses.echo()
-    curses.endwin()
-
-
 win = curses.initscr()
-atexit.register(tear_down)
-curses.endwin()
 lineno = 0
+colors_map = dict(
+    green=3,
+    red=10,
+    yellow=4,
+)
 
 
-def print_line(line, highlight=False):
+def printl(line, color=None, bold=False, highlight=False):
     """A thin wrapper around curses's addstr()."""
     global lineno
     try:
+        flags = 0
+        if color:
+            flags |= curses.color_pair(colors_map[color])
+        if bold:
+            flags |= curses.A_BOLD
         if highlight:
             line += " " * (win.getmaxyx()[1] - len(line))
-            win.addstr(lineno, 0, line, curses.A_REVERSE)
-        else:
-            win.addstr(lineno, 0, line, 0)
+            flags |= curses.A_STANDOUT
+        win.addstr(lineno, 0, line, flags)
     except curses.error:
         lineno = 0
         win.refresh()
@@ -105,6 +102,15 @@ def poll(interval):
     return (processes, procs_status)
 
 
+def get_color(perc):
+    if perc <= 30:
+        return "green"
+    elif perc <= 80:
+        return "yellow"
+    else:
+        return "red"
+
+
 def print_header(procs_status, num_procs):
     """Print system-related info, above the process list."""
 
@@ -117,8 +123,8 @@ def print_header(procs_status, num_procs):
     percs = psutil.cpu_percent(interval=0, percpu=True)
     for cpu_num, perc in enumerate(percs):
         dashes, empty_dashes = get_dashes(perc)
-        print_line(" CPU%-2s [%s%s] %5s%%" % (cpu_num, dashes, empty_dashes,
-                                              perc))
+        line = " CPU%-2s [%s%s] %5s%%" % (cpu_num, dashes, empty_dashes, perc)
+        printl(line, color=get_color(perc))
     mem = psutil.virtual_memory()
     dashes, empty_dashes = get_dashes(mem.percent)
     line = " Mem   [%s%s] %5s%% %6s / %s" % (
@@ -127,7 +133,7 @@ def print_header(procs_status, num_procs):
         str(int(mem.used / 1024 / 1024)) + "M",
         str(int(mem.total / 1024 / 1024)) + "M"
     )
-    print_line(line)
+    printl(line, color=get_color(mem.percent))
 
     # swap usage
     swap = psutil.swap_memory()
@@ -138,7 +144,7 @@ def print_header(procs_status, num_procs):
         str(int(swap.used / 1024 / 1024)) + "M",
         str(int(swap.total / 1024 / 1024)) + "M"
     )
-    print_line(line)
+    printl(line, color=get_color(swap.percent))
 
     # processes number and status
     st = []
@@ -146,14 +152,14 @@ def print_header(procs_status, num_procs):
         if y:
             st.append("%s=%s" % (x, y))
     st.sort(key=lambda x: x[:3] in ('run', 'sle'), reverse=1)
-    print_line(" Processes: %s (%s)" % (num_procs, ', '.join(st)))
+    printl(" Processes: %s (%s)" % (num_procs, ', '.join(st)))
     # load average, uptime
     uptime = datetime.datetime.now() - \
         datetime.datetime.fromtimestamp(psutil.boot_time())
     av1, av2, av3 = psutil.getloadavg()
     line = " Load average: %.2f %.2f %.2f  Uptime: %s" \
         % (av1, av2, av3, str(uptime).split('.')[0])
-    print_line(line)
+    printl(line)
 
 
 def refresh_window(procs, procs_status):
@@ -164,8 +170,8 @@ def refresh_window(procs, procs_status):
     header = templ % ("PID", "USER", "NI", "VIRT", "RES", "CPU%", "MEM%",
                       "TIME+", "NAME")
     print_header(procs_status, len(procs))
-    print_line("")
-    print_line(header, highlight=True)
+    printl("")
+    printl(header, bold=True, highlight=True)
     for p in procs:
         # TIME+ column shows process CPU cumulative time and it
         # is expressed as: "mm:ss.ms"
@@ -197,21 +203,42 @@ def refresh_window(procs, procs_status):
                         p.dict['name'] or '',
                         )
         try:
-            print_line(line)
+            printl(line)
         except curses.error:
             break
         win.refresh()
 
 
+def setup():
+    curses.start_color()
+    curses.use_default_colors()
+    for i in range(0, curses.COLORS):
+        curses.init_pair(i + 1, i, -1)
+    curses.endwin()
+    win.nodelay(1)
+
+
+def tear_down():
+    win.keypad(0)
+    curses.nocbreak()
+    curses.echo()
+    curses.endwin()
+
+
 def main():
+    setup()
     try:
         interval = 0
         while True:
+            if win.getch() == ord('q'):
+                break
             args = poll(interval)
             refresh_window(*args)
             interval = 1
     except (KeyboardInterrupt, SystemExit):
         pass
+    finally:
+        tear_down()
 
 
 if __name__ == '__main__':