"""Misc. useful functionality used by the rest of this package.
This module provides common functionality used by the other modules in
this package.
"""
import sys
import os
import subprocess
try:
from subprocess import CalledProcessError
except ImportError:
class CalledProcessError(Exception):
"""This exception is raised when a process run by check_call() returns
a non-zero exit status. The exit status will be stored in the
returncode attribute."""
def __init__(self, returncode, cmd):
self.returncode = returncode
self.cmd = cmd
def __str__(self):
return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)
DEBUG = False
def notify(msg, *args):
"""Print a message to stderr."""
print >> sys.stderr, msg % args
def debug (msg, *args):
"""Print a debug message to stderr when DEBUG is enabled."""
if DEBUG:
print >> sys.stderr, msg % args
def error (msg, *args):
"""Print an error message to stderr."""
print >> sys.stderr, "ERROR:", msg % args
def warn(msg, *args):
"""Print a warning message to stderr."""
print >> sys.stderr, "warning:", msg % args
def die (msg, *args):
"""Print as error message to stderr and exit the program."""
error(msg, *args)
sys.exit(1)
class ProgressIndicator(object):
"""Simple progress indicator.
Displayed as a spinning character by default, but can be customized
by passing custom messages that overrides the spinning character.
"""
States = ("|", "/", "-", "\\")
def __init__ (self, prefix = "", f = sys.stdout):
"""Create a new ProgressIndicator, bound to the given file object."""
self.n = 0 self.f = f self.prev_len = 0 self.prefix = prefix self.prefix_lens = []
def pushprefix (self, prefix):
"""Append the given prefix onto the prefix stack."""
self.prefix_lens.append(len(self.prefix))
self.prefix += prefix
def popprefix (self):
"""Remove the last prefix from the prefix stack."""
prev_len = self.prefix_lens.pop()
self.prefix = self.prefix[:prev_len]
def __call__ (self, msg = None, lf = False):
"""Indicate progress, possibly with a custom message."""
if msg is None:
msg = self.States[self.n % len(self.States)]
msg = self.prefix + msg
print >> self.f, "\r%-*s" % (self.prev_len, msg),
self.prev_len = len(msg.expandtabs())
if lf:
print >> self.f
self.prev_len = 0
self.n += 1
def finish (self, msg = "done", noprefix = False):
"""Finalize progress indication with the given message."""
if noprefix:
self.prefix = ""
self(msg, True)
def start_command (args, cwd = None, shell = False, add_env = None,
stdin = subprocess.PIPE, stdout = subprocess.PIPE,
stderr = subprocess.PIPE):
"""Start the given command, and return a subprocess object.
This provides a simpler interface to the subprocess module.
"""
env = None
if add_env is not None:
env = os.environ.copy()
env.update(add_env)
return subprocess.Popen(args, bufsize = 1, stdin = stdin, stdout = stdout,
stderr = stderr, cwd = cwd, shell = shell,
env = env, universal_newlines = True)
def run_command (args, cwd = None, shell = False, add_env = None,
flag_error = True):
"""Run the given command to completion, and return its results.
This provides a simpler interface to the subprocess module.
The results are formatted as a 3-tuple: (exit_code, output, errors)
If flag_error is enabled, Error messages will be produced if the
subprocess terminated with a non-zero exit code and/or stderr
output.
The other arguments are passed on to start_command().
"""
process = start_command(args, cwd, shell, add_env)
(output, errors) = process.communicate()
exit_code = process.returncode
if flag_error and errors:
error("'%s' returned errors:\n---\n%s---", " ".join(args), errors)
if flag_error and exit_code:
error("'%s' returned exit code %i", " ".join(args), exit_code)
return (exit_code, output, errors)
def call(*popenargs, **kwargs):
"""Run command with arguments. Wait for command to complete, then
return the returncode attribute.
The arguments are the same as for the Popen constructor. Example:
retcode = call(["ls", "-l"])
"""
return subprocess.Popen(*popenargs, **kwargs).wait()
def check_call(*popenargs, **kwargs):
"""Run command with arguments. Wait for command to complete. If
the exit code was zero then return, otherwise raise
CalledProcessError. The CalledProcessError object will have the
return code in the returncode attribute.
The arguments are the same as for the Popen constructor. Example:
check_call(["ls", "-l"])
"""
retcode = call(*popenargs, **kwargs)
if retcode:
cmd = kwargs.get("args")
if cmd is None:
cmd = popenargs[0]
raise CalledProcessError(retcode, cmd)
return 0
def check_output(*popenargs, **kwargs):
r"""Run command with arguments and return its output as a byte string.
If the exit code was non-zero it raises a CalledProcessError. The
CalledProcessError object will have the return code in the returncode
attribute and output in the output attribute.
The arguments are the same as for the Popen constructor. Example:
>>> check_output(["ls", "-l", "/dev/null"])
'crw-rw-rw- 1 root root 1, 3 Oct 18 2007 /dev/null\n'
The stdout argument is not allowed as it is used internally.
To capture standard error in the result, use stderr=STDOUT.
>>> check_output(["/bin/sh", "-c",
... "ls -l non_existent_file ; exit 0"],
... stderr=STDOUT)
'ls: non_existent_file: No such file or directory\n'
"""
if 'stdout' in kwargs:
raise ValueError('stdout argument not allowed, it will be overridden.')
process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
output, unused_err = process.communicate()
retcode = process.poll()
if retcode:
cmd = kwargs.get("args")
if cmd is None:
cmd = popenargs[0]
raise subprocess.CalledProcessError(retcode, cmd)
return output
def file_reader_method (missing_ok = False):
"""Decorator for simplifying reading of files.
If missing_ok is True, a failure to open a file for reading will
not raise the usual IOError, but instead the wrapped method will be
called with f == None. The method must in this case properly
handle f == None.
"""
def _wrap (method):
"""Teach given method to handle both filenames and file objects.
The given method must take a file object as its second argument
(the first argument being 'self', of course). This decorator
will take a filename given as the second argument and promote
it to a file object.
"""
def _wrapped_method (self, filename, *args, **kwargs):
if isinstance(filename, file):
f = filename
else:
try:
f = open(filename, 'r')
except IOError:
if missing_ok:
f = None
else:
raise
try:
return method(self, f, *args, **kwargs)
finally:
if not isinstance(filename, file) and f:
f.close()
return _wrapped_method
return _wrap
def file_writer_method (method):
"""Decorator for simplifying writing of files.
Enables the given method to handle both filenames and file objects.
The given method must take a file object as its second argument
(the first argument being 'self', of course). This decorator will
take a filename given as the second argument and promote it to a
file object.
"""
def _new_method (self, filename, *args, **kwargs):
if isinstance(filename, file):
f = filename
else:
parent_dir = os.path.dirname(filename)
if not os.path.isdir(parent_dir):
os.makedirs(parent_dir)
f = open(filename, 'w')
try:
return method(self, f, *args, **kwargs)
finally:
if not isinstance(filename, file):
f.close()
return _new_method