git_remote_helpers / util.pyon commit Merge branch 'maint' (2f80de9)
   1#!/usr/bin/env python
   2
   3"""Misc. useful functionality used by the rest of this package.
   4
   5This module provides common functionality used by the other modules in
   6this package.
   7
   8"""
   9
  10import sys
  11import os
  12import subprocess
  13
  14
  15# Whether or not to show debug messages
  16DEBUG = False
  17
  18def notify(msg, *args):
  19    """Print a message to stderr."""
  20    print >> sys.stderr, msg % args
  21
  22def debug (msg, *args):
  23    """Print a debug message to stderr when DEBUG is enabled."""
  24    if DEBUG:
  25        print >> sys.stderr, msg % args
  26
  27def error (msg, *args):
  28    """Print an error message to stderr."""
  29    print >> sys.stderr, "ERROR:", msg % args
  30
  31def warn(msg, *args):
  32    """Print a warning message to stderr."""
  33    print >> sys.stderr, "warning:", msg % args
  34
  35def die (msg, *args):
  36    """Print as error message to stderr and exit the program."""
  37    error(msg, *args)
  38    sys.exit(1)
  39
  40
  41class ProgressIndicator(object):
  42
  43    """Simple progress indicator.
  44
  45    Displayed as a spinning character by default, but can be customized
  46    by passing custom messages that overrides the spinning character.
  47
  48    """
  49
  50    States = ("|", "/", "-", "\\")
  51
  52    def __init__ (self, prefix = "", f = sys.stdout):
  53        """Create a new ProgressIndicator, bound to the given file object."""
  54        self.n = 0  # Simple progress counter
  55        self.f = f  # Progress is written to this file object
  56        self.prev_len = 0  # Length of previous msg (to be overwritten)
  57        self.prefix = prefix  # Prefix prepended to each progress message
  58        self.prefix_lens = [] # Stack of prefix string lengths
  59
  60    def pushprefix (self, prefix):
  61        """Append the given prefix onto the prefix stack."""
  62        self.prefix_lens.append(len(self.prefix))
  63        self.prefix += prefix
  64
  65    def popprefix (self):
  66        """Remove the last prefix from the prefix stack."""
  67        prev_len = self.prefix_lens.pop()
  68        self.prefix = self.prefix[:prev_len]
  69
  70    def __call__ (self, msg = None, lf = False):
  71        """Indicate progress, possibly with a custom message."""
  72        if msg is None:
  73            msg = self.States[self.n % len(self.States)]
  74        msg = self.prefix + msg
  75        print >> self.f, "\r%-*s" % (self.prev_len, msg),
  76        self.prev_len = len(msg.expandtabs())
  77        if lf:
  78            print >> self.f
  79            self.prev_len = 0
  80        self.n += 1
  81
  82    def finish (self, msg = "done", noprefix = False):
  83        """Finalize progress indication with the given message."""
  84        if noprefix:
  85            self.prefix = ""
  86        self(msg, True)
  87
  88
  89def start_command (args, cwd = None, shell = False, add_env = None,
  90                   stdin = subprocess.PIPE, stdout = subprocess.PIPE,
  91                   stderr = subprocess.PIPE):
  92    """Start the given command, and return a subprocess object.
  93
  94    This provides a simpler interface to the subprocess module.
  95
  96    """
  97    env = None
  98    if add_env is not None:
  99        env = os.environ.copy()
 100        env.update(add_env)
 101    return subprocess.Popen(args, bufsize = 1, stdin = stdin, stdout = stdout,
 102                            stderr = stderr, cwd = cwd, shell = shell,
 103                            env = env, universal_newlines = True)
 104
 105
 106def run_command (args, cwd = None, shell = False, add_env = None,
 107                 flag_error = True):
 108    """Run the given command to completion, and return its results.
 109
 110    This provides a simpler interface to the subprocess module.
 111
 112    The results are formatted as a 3-tuple: (exit_code, output, errors)
 113
 114    If flag_error is enabled, Error messages will be produced if the
 115    subprocess terminated with a non-zero exit code and/or stderr
 116    output.
 117
 118    The other arguments are passed on to start_command().
 119
 120    """
 121    process = start_command(args, cwd, shell, add_env)
 122    (output, errors) = process.communicate()
 123    exit_code = process.returncode
 124    if flag_error and errors:
 125        error("'%s' returned errors:\n---\n%s---", " ".join(args), errors)
 126    if flag_error and exit_code:
 127        error("'%s' returned exit code %i", " ".join(args), exit_code)
 128    return (exit_code, output, errors)
 129
 130
 131def file_reader_method (missing_ok = False):
 132    """Decorator for simplifying reading of files.
 133
 134    If missing_ok is True, a failure to open a file for reading will
 135    not raise the usual IOError, but instead the wrapped method will be
 136    called with f == None.  The method must in this case properly
 137    handle f == None.
 138
 139    """
 140    def _wrap (method):
 141        """Teach given method to handle both filenames and file objects.
 142
 143        The given method must take a file object as its second argument
 144        (the first argument being 'self', of course).  This decorator
 145        will take a filename given as the second argument and promote
 146        it to a file object.
 147
 148        """
 149        def _wrapped_method (self, filename, *args, **kwargs):
 150            if isinstance(filename, file):
 151                f = filename
 152            else:
 153                try:
 154                    f = open(filename, 'r')
 155                except IOError:
 156                    if missing_ok:
 157                        f = None
 158                    else:
 159                        raise
 160            try:
 161                return method(self, f, *args, **kwargs)
 162            finally:
 163                if not isinstance(filename, file) and f:
 164                    f.close()
 165        return _wrapped_method
 166    return _wrap
 167
 168
 169def file_writer_method (method):
 170    """Decorator for simplifying writing of files.
 171
 172    Enables the given method to handle both filenames and file objects.
 173
 174    The given method must take a file object as its second argument
 175    (the first argument being 'self', of course).  This decorator will
 176    take a filename given as the second argument and promote it to a
 177    file object.
 178
 179    """
 180    def _new_method (self, filename, *args, **kwargs):
 181        if isinstance(filename, file):
 182            f = filename
 183        else:
 184            # Make sure the containing directory exists
 185            parent_dir = os.path.dirname(filename)
 186            if not os.path.isdir(parent_dir):
 187                os.makedirs(parent_dir)
 188            f = open(filename, 'w')
 189        try:
 190            return method(self, f, *args, **kwargs)
 191        finally:
 192            if not isinstance(filename, file):
 193                f.close()
 194    return _new_method