| #!/usr/bin/env python |
| |
| """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 |
| |
| |
| # Whether or not to show debug messages |
| 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 # Simple progress counter |
| self.f = f # Progress is written to this file object |
| self.prev_len = 0 # Length of previous msg (to be overwritten) |
| self.prefix = prefix # Prefix prepended to each progress message |
| self.prefix_lens = [] # Stack of prefix string lengths |
| |
| 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 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: |
| # Make sure the containing directory exists |
| 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 |
