michael@0: #!/usr/bin/env python
michael@0: 
michael@0: # Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
michael@0: # Use of this source code is governed by a BSD-style license that can be
michael@0: # found in the LICENSE file.
michael@0: 
michael@0: """
michael@0: A clone of iotop (http://guichaz.free.fr/iotop/) showing real time
michael@0: disk I/O statistics.
michael@0: 
michael@0: It works on Linux only (FreeBSD and OSX are missing support for IO
michael@0: counters).
michael@0: It doesn't work on Windows as curses module is required.
michael@0: 
michael@0: Author: Giampaolo Rodola' <g.rodola@gmail.com>
michael@0: """
michael@0: 
michael@0: import os
michael@0: import sys
michael@0: import psutil
michael@0: if not hasattr(psutil.Process, 'get_io_counters') or os.name != 'posix':
michael@0:     sys.exit('platform not supported')
michael@0: import time
michael@0: import curses
michael@0: import atexit
michael@0: 
michael@0: 
michael@0: # --- curses stuff
michael@0: def tear_down():
michael@0:     win.keypad(0)
michael@0:     curses.nocbreak()
michael@0:     curses.echo()
michael@0:     curses.endwin()
michael@0: 
michael@0: win = curses.initscr()
michael@0: atexit.register(tear_down)
michael@0: curses.endwin()
michael@0: lineno = 0
michael@0: 
michael@0: def print_line(line, highlight=False):
michael@0:     """A thin wrapper around curses's addstr()."""
michael@0:     global lineno
michael@0:     try:
michael@0:         if highlight:
michael@0:             line += " " * (win.getmaxyx()[1] - len(line))
michael@0:             win.addstr(lineno, 0, line, curses.A_REVERSE)
michael@0:         else:
michael@0:             win.addstr(lineno, 0, line, 0)
michael@0:     except curses.error:
michael@0:         lineno = 0
michael@0:         win.refresh()
michael@0:         raise
michael@0:     else:
michael@0:         lineno += 1
michael@0: # --- /curses stuff
michael@0: 
michael@0: 
michael@0: def bytes2human(n):
michael@0:     """
michael@0:     >>> bytes2human(10000)
michael@0:     '9.8 K/s'
michael@0:     >>> bytes2human(100001221)
michael@0:     '95.4 M/s'
michael@0:     """
michael@0:     symbols = ('K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
michael@0:     prefix = {}
michael@0:     for i, s in enumerate(symbols):
michael@0:         prefix[s] = 1 << (i+1)*10
michael@0:     for s in reversed(symbols):
michael@0:         if n >= prefix[s]:
michael@0:             value = float(n) / prefix[s]
michael@0:             return '%.2f %s/s' % (value, s)
michael@0:     return '%.2f B/s' % (n)
michael@0: 
michael@0: def poll(interval):
michael@0:     """Calculate IO usage by comparing IO statics before and
michael@0:     after the interval.
michael@0:     Return a tuple including all currently running processes
michael@0:     sorted by IO activity and total disks I/O activity.
michael@0:     """
michael@0:     # first get a list of all processes and disk io counters
michael@0:     procs = [p for p in psutil.process_iter()]
michael@0:     for p in procs[:]:
michael@0:         try:
michael@0:             p._before = p.get_io_counters()
michael@0:         except psutil.Error:
michael@0:             procs.remove(p)
michael@0:             continue
michael@0:     disks_before = psutil.disk_io_counters()
michael@0: 
michael@0:     # sleep some time
michael@0:     time.sleep(interval)
michael@0: 
michael@0:     # then retrieve the same info again
michael@0:     for p in procs[:]:
michael@0:         try:
michael@0:             p._after = p.get_io_counters()
michael@0:             p._cmdline = ' '.join(p.cmdline)
michael@0:             if not p._cmdline:
michael@0:                 p._cmdline = p.name
michael@0:             p._username = p.username
michael@0:         except psutil.NoSuchProcess:
michael@0:             procs.remove(p)
michael@0:     disks_after = psutil.disk_io_counters()
michael@0: 
michael@0:     # finally calculate results by comparing data before and
michael@0:     # after the interval
michael@0:     for p in procs:
michael@0:         p._read_per_sec = p._after.read_bytes - p._before.read_bytes
michael@0:         p._write_per_sec = p._after.write_bytes - p._before.write_bytes
michael@0:         p._total = p._read_per_sec + p._write_per_sec
michael@0: 
michael@0:     disks_read_per_sec = disks_after.read_bytes - disks_before.read_bytes
michael@0:     disks_write_per_sec = disks_after.write_bytes - disks_before.write_bytes
michael@0: 
michael@0:     # sort processes by total disk IO so that the more intensive
michael@0:     # ones get listed first
michael@0:     processes = sorted(procs, key=lambda p: p._total, reverse=True)
michael@0: 
michael@0:     return (processes, disks_read_per_sec, disks_write_per_sec)
michael@0: 
michael@0: 
michael@0: def refresh_window(procs, disks_read, disks_write):
michael@0:     """Print results on screen by using curses."""
michael@0:     curses.endwin()
michael@0:     templ = "%-5s %-7s %11s %11s  %s"
michael@0:     win.erase()
michael@0: 
michael@0:     disks_tot = "Total DISK READ: %s | Total DISK WRITE: %s" \
michael@0:                 % (bytes2human(disks_read), bytes2human(disks_write))
michael@0:     print_line(disks_tot)
michael@0: 
michael@0:     header = templ % ("PID", "USER", "DISK READ", "DISK WRITE", "COMMAND")
michael@0:     print_line(header, highlight=True)
michael@0: 
michael@0:     for p in procs:
michael@0:         line = templ % (p.pid,
michael@0:                         p._username[:7],
michael@0:                         bytes2human(p._read_per_sec),
michael@0:                         bytes2human(p._write_per_sec),
michael@0:                         p._cmdline)
michael@0:         try:
michael@0:             print_line(line)
michael@0:         except curses.error:
michael@0:             break
michael@0:     win.refresh()
michael@0: 
michael@0: def main():
michael@0:     try:
michael@0:         interval = 0
michael@0:         while 1:
michael@0:             args = poll(interval)
michael@0:             refresh_window(*args)
michael@0:             interval = 1
michael@0:     except (KeyboardInterrupt, SystemExit):
michael@0:         pass
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     main()