git_remote_helpers/util.py

   1 #!/usr/bin/env python
   2
   3 """Misc. useful functionality used by the rest of this package.
   4
   5 This module provides common functionality used by the other modules in
   6 this package.
   7
   8 """
   9
  10 import sys
  11 import os
  12 import subprocess
  13
  14 try:
  15     from subprocess import CalledProcessError
  16 except ImportError:
  17     # from python2.7:subprocess.py
  18     # Exception classes used by this module.
  19     class CalledProcessError(Exception):
  20         """This exception is raised when a process run by check_call() returns
  21         a non-zero exit status.  The exit status will be stored in the
  22         returncode attribute."""
  23         def __init__(self, returncode, cmd):
  24             self.returncode = returncode
  25             self.cmd = cmd
  26         def __str__(self):
  27             return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)
  28
  29
  30 # Whether or not to show debug messages
  31 DEBUG = False
  32
  33 def notify(msg, *args):
  34     """Print a message to stderr."""
  35     print >> sys.stderr, msg % args
  36
  37 def debug (msg, *args):
  38     """Print a debug message to stderr when DEBUG is enabled."""
  39     if DEBUG:
  40         print >> sys.stderr, msg % args
  41
  42 def error (msg, *args):
  43     """Print an error message to stderr."""
  44     print >> sys.stderr, "ERROR:", msg % args
  45
  46 def warn(msg, *args):
  47     """Print a warning message to stderr."""
  48     print >> sys.stderr, "warning:", msg % args
  49
  50 def die (msg, *args):
  51     """Print as error message to stderr and exit the program."""
  52     error(msg, *args)
  53     sys.exit(1)
  54
  55
  56 class ProgressIndicator(object):
  57
  58     """Simple progress indicator.
  59
  60     Displayed as a spinning character by default, but can be customized
  61     by passing custom messages that overrides the spinning character.
  62
  63     """
  64
  65     States = ("|", "/", "-", "\\")
  66
  67     def __init__ (self, prefix = "", f = sys.stdout):
  68         """Create a new ProgressIndicator, bound to the given file object."""
  69         self.n = 0  # Simple progress counter
  70         self.f = f  # Progress is written to this file object
  71         self.prev_len = 0  # Length of previous msg (to be overwritten)
  72         self.prefix = prefix  # Prefix prepended to each progress message
  73         self.prefix_lens = [] # Stack of prefix string lengths
  74
  75     def pushprefix (self, prefix):
  76         """Append the given prefix onto the prefix stack."""
  77         self.prefix_lens.append(len(self.prefix))
  78         self.prefix += prefix
  79
  80     def popprefix (self):
  81         """Remove the last prefix from the prefix stack."""
  82         prev_len = self.prefix_lens.pop()
  83         self.prefix = self.prefix[:prev_len]
  84
  85     def __call__ (self, msg = None, lf = False):
  86         """Indicate progress, possibly with a custom message."""
  87         if msg is None:
  88             msg = self.States[self.n % len(self.States)]
  89         msg = self.prefix + msg
  90         print >> self.f, "\r%-*s" % (self.prev_len, msg),
  91         self.prev_len = len(msg.expandtabs())
  92         if lf:
  93             print >> self.f
  94             self.prev_len = 0
  95         self.n += 1
  96
  97     def finish (self, msg = "done", noprefix = False):
  98         """Finalize progress indication with the given message."""
  99         if noprefix:
 100             self.prefix = ""
 101         self(msg, True)
 102
 103
 104 def start_command (args, cwd = None, shell = False, add_env = None,
 105                    stdin = subprocess.PIPE, stdout = subprocess.PIPE,
 106                    stderr = subprocess.PIPE):
 107     """Start the given command, and return a subprocess object.
 108
 109     This provides a simpler interface to the subprocess module.
 110
 111     """
 112     env = None
 113     if add_env is not None:
 114         env = os.environ.copy()
 115         env.update(add_env)
 116     return subprocess.Popen(args, bufsize = 1, stdin = stdin, stdout = stdout,
 117                             stderr = stderr, cwd = cwd, shell = shell,
 118                             env = env, universal_newlines = True)
 119
 120
 121 def run_command (args, cwd = None, shell = False, add_env = None,
 122                  flag_error = True):
 123     """Run the given command to completion, and return its results.
 124
 125     This provides a simpler interface to the subprocess module.
 126
 127     The results are formatted as a 3-tuple: (exit_code, output, errors)
 128
 129     If flag_error is enabled, Error messages will be produced if the
 130     subprocess terminated with a non-zero exit code and/or stderr
 131     output.
 132
 133     The other arguments are passed on to start_command().
 134
 135     """
 136     process = start_command(args, cwd, shell, add_env)
 137     (output, errors) = process.communicate()
 138     exit_code = process.returncode
 139     if flag_error and errors:
 140         error("'%s' returned errors:\n---\n%s---", " ".join(args), errors)
 141     if flag_error and exit_code:
 142         error("'%s' returned exit code %i", " ".join(args), exit_code)
 143     return (exit_code, output, errors)
 144
 145
 146 # from python2.7:subprocess.py
 147 def call(*popenargs, **kwargs):
 148     """Run command with arguments.  Wait for command to complete, then
 149     return the returncode attribute.
 150
 151     The arguments are the same as for the Popen constructor.  Example:
 152
 153     retcode = call(["ls", "-l"])
 154     """
 155     return subprocess.Popen(*popenargs, **kwargs).wait()
 156
 157
 158 # from python2.7:subprocess.py
 159 def check_call(*popenargs, **kwargs):
 160     """Run command with arguments.  Wait for command to complete.  If
 161     the exit code was zero then return, otherwise raise
 162     CalledProcessError.  The CalledProcessError object will have the
 163     return code in the returncode attribute.
 164
 165     The arguments are the same as for the Popen constructor.  Example:
 166
 167     check_call(["ls", "-l"])
 168     """
 169     retcode = call(*popenargs, **kwargs)
 170     if retcode:
 171         cmd = kwargs.get("args")
 172         if cmd is None:
 173             cmd = popenargs[0]
 174         raise CalledProcessError(retcode, cmd)
 175     return 0
 176
 177
 178 # from python2.7:subprocess.py
 179 def check_output(*popenargs, **kwargs):
 180     r"""Run command with arguments and return its output as a byte string.
 181
 182     If the exit code was non-zero it raises a CalledProcessError.  The
 183     CalledProcessError object will have the return code in the returncode
 184     attribute and output in the output attribute.
 185
 186     The arguments are the same as for the Popen constructor.  Example:
 187
 188     >>> check_output(["ls", "-l", "/dev/null"])
 189     'crw-rw-rw- 1 root root 1, 3 Oct 18  2007 /dev/null\n'
 190
 191     The stdout argument is not allowed as it is used internally.
 192     To capture standard error in the result, use stderr=STDOUT.
 193
 194     >>> check_output(["/bin/sh", "-c",
 195     ...               "ls -l non_existent_file ; exit 0"],
 196     ...              stderr=STDOUT)
 197     'ls: non_existent_file: No such file or directory\n'
 198     """
 199     if 'stdout' in kwargs:
 200         raise ValueError('stdout argument not allowed, it will be overridden.')
 201     process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
 202     output, unused_err = process.communicate()
 203     retcode = process.poll()
 204     if retcode:
 205         cmd = kwargs.get("args")
 206         if cmd is None:
 207             cmd = popenargs[0]
 208         raise subprocess.CalledProcessError(retcode, cmd)
 209     return output
 210
 211
 212 def file_reader_method (missing_ok = False):
 213     """Decorator for simplifying reading of files.
 214
 215     If missing_ok is True, a failure to open a file for reading will
 216     not raise the usual IOError, but instead the wrapped method will be
 217     called with f == None.  The method must in this case properly
 218     handle f == None.
 219
 220     """
 221     def _wrap (method):
 222         """Teach given method to handle both filenames and file objects.
 223
 224         The given method must take a file object as its second argument
 225         (the first argument being 'self', of course).  This decorator
 226         will take a filename given as the second argument and promote
 227         it to a file object.
 228
 229         """
 230         def _wrapped_method (self, filename, *args, **kwargs):
 231             if isinstance(filename, file):
 232                 f = filename
 233             else:
 234                 try:
 235                     f = open(filename, 'r')
 236                 except IOError:
 237                     if missing_ok:
 238                         f = None
 239                     else:
 240                         raise
 241             try:
 242                 return method(self, f, *args, **kwargs)
 243             finally:
 244                 if not isinstance(filename, file) and f:
 245                     f.close()
 246         return _wrapped_method
 247     return _wrap
 248
 249
 250 def file_writer_method (method):
 251     """Decorator for simplifying writing of files.
 252
 253     Enables the given method to handle both filenames and file objects.
 254
 255     The given method must take a file object as its second argument
 256     (the first argument being 'self', of course).  This decorator will
 257     take a filename given as the second argument and promote it to a
 258     file object.
 259
 260     """
 261     def _new_method (self, filename, *args, **kwargs):
 262         if isinstance(filename, file):
 263             f = filename
 264         else:
 265             # Make sure the containing directory exists
 266             parent_dir = os.path.dirname(filename)
 267             if not os.path.isdir(parent_dir):
 268                 os.makedirs(parent_dir)
 269             f = open(filename, 'w')
 270         try:
 271             return method(self, f, *args, **kwargs)
 272         finally:
 273             if not isinstance(filename, file):
 274                 f.close()
 275     return _new_method