Python module that runs stat collectors (like iostat, mpstat) on background

bench.py

import os, subprocess
import errno
import time
import shutil
import signal

from string import maketrans
from datetime import datetime
from threading import Thread
	
def arg_mangle(args):
    """Replaces '/' to '@' and removes dots and spaces
    from list of args. This is needed to create correct
    output filenames, e.g:
    cp /some/file . -> cp@some@file """
    return ''.join([arg.translate(maketrans('/', '@'), '. ') 
		    for arg 
		    in args])

class command_not_found:
    def __init__(self, cmdname, cmdpath):
	self.cmdname = cmdname
	self.cmdpath = cmdpath
	
    def __str__(self):
	return "%s not found in %s" % (self.cmdname, ', '.join(self.cmdpath))

def cmdsearch(cmd, cmdpathlist):
    """Searches command in cmdpath.
	If not found returns command_not_found exception"""
    for cmdpath in cmdpathlist:
	fullcmd = os.path.join(cmdpath, cmd)  
	if os.path.exists(fullcmd) and os.access(fullcmd, os.X_OK):
	    return fullcmd;
    else:
	raise command_not_found(cmd, statistic.cmdpath)

class statcmd:
    cmdpathlist = ['/usr/bin', '/bin', '/sbin', '/usr/sbin', '/usr/local/bin']
	
    def __init__(self, cmd, args, interval = 5, runonce = True, filename = ''):
	self.statcmd = cmdsearch(cmd, statcmd.cmdpathlist)
	if not filename:
	    self.filename = cmd + arg_mangle(args)
	    if len(self.filename) > 24:
		self.filename = self.filename[:24]
	else:
	    self.filename = filename
	self.cmdargs = args
	self.interval = interval
	self.runonce = runonce
	
    def setoutdir(self, outdir):
	self.statfile = os.path.join(outdir, self.filename)
	self.statout = file(self.statfile + '.out', 'w') 
	self.staterr = file(self.statfile + '.err', 'w') 
    
    def getinfo(self):
	return (self.interval, self.runonce)
    
    def start(self):
	"""Runs process"""
	self.proc = subprocess.Popen([self.statcmd] + self.cmdargs, shell=False, 
		stdout=self.statout, stderr=self.staterr)
	
    def poll(self):
	""""Polls is collector alive?"""
	# TODO: do real poll :)
	return 0
    
    def stop(self):
	"""Stops process if exists using SIGINT SIGTERM"""
	try:
	    os.kill(self.proc.pid, signal.SIGINT)
	    time.sleep(2);
	    os.kill(self.proc.pid, signal.SIGTERM)
	    
	except OSError, err:
	    if err.errno != errno.ESRCH:
		raise
	    
    def __str__(self):
	return "<%s [cmd: %s, out: %s.*]>" % (self.__class__.__name__, 
		    ' '.join([self.statcmd] + self.cmdargs), self.statfile) 

class statcollector(Thread):
    """ Collector is thread object which proxies statobjs
	It is helper class, use statcmd or your own classes
	
	Collector is a proxy for statobj, which can be statcmd object
	(runs system command) or any other python object which implements:
	    - setoutdir(outdir) - sets output directory for statobj
	    - getinfo() - returns tuple of thread behaviour (interval, runonce)
	    - start() - starts stat collection
	    - stop() - stops it
	    - __str__() 
	    
	if runonce is True, statobj.start() calls every interval settings
	othervise statobj.start() calls on run() and statobj.stop() on done"""
    
    def __init__(self, statobj, outdir):
	Thread.__init__(self)
	
	self.obj = statobj;
	self.obj.setoutdir(outdir)	
	(self.interval, self.runonce) = self.obj.getinfo();
	
	self.isrun = 1;

    def run(self):
	""" Main thread loop"""
	if self.runonce:
	    self.obj.start();
	    
	while self.isrun:
	    if self.runonce:
		if not self.obj.poll():
		    # Collector dead - restart
		    self.obj.start();
		time.sleep(self.interval);
	    else:
		self.obj.start();
		time.sleep(self.interval);
		self.obj.stop();
    
    def done(self):
	"""Stops thread and stops proxy object"""
	if self.runonce:
	    self.obj.stop();
	self.isrun = 0;
	
    def __str__(self):
	return "statcollector proxy for %s" % self.statobj
	
      
class test_failure:
    def __init__(self, cmd, code):
	self.cmd = cmd
	self.code = code
	    
    def __str__(self):
	print "Command %s returned %d code" % self.code
	
def run_test(cmd, stats, outfile, outdir = '/tmp'):
    """Runs command cmd and stats all output are gathered into outdir
    For information on stats see statcollector and statcmd"""
    os.umask(0)
    os.makedirs(outdir, 0777);
	
    try:
	filename = os.path.join(outdir, outfile)
	print "Running tests for %s" % cmd
	
	statthreads = []
		
	# Runs all statistics in individual threads
	for stat in stats:
	    thread = statcollector(stat, outdir)
	    statthreads.append(thread)
	    thread.start()
	    print "Running %s" % stat
	
	try:
	    cmd = cmd + ' 1>%s.out 2>%s.err' % (filename, filename)
	    print "Running %s" % cmd
	    testret = subprocess.call(cmd, shell = True)
	    
	    if testret != 0:
		raise test_failure(cmd, testret)
	finally:
	    # Stop gathering statistics
	    for thread in statthreads:
		thread.done()	
    except:
	# In case of error recursively remove all output
	print 'Exiting'
	
	answer = ''
	while not answer or answer not in 'yYnN':
	    answer = raw_input('Clean output directory? [y/n]: ')
	    if answer in 'nN':
		raise
	
	shutil.rmtree(outdir)
	raise

    print "Finished"