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
  15 # Whether or not to show debug messages
  16 DEBUG = False
  17
  18 def notify(msg, *args):
  19     """Print a message to stderr."""
  20     print >> sys.stderr, msg % args
  21
  22 def debug (msg, *args):
  23     """Print a debug message to stderr when DEBUG is enabled."""
  24     if DEBUG:
  25         print >> sys.stderr, msg % args
  26
  27 def error (msg, *args):
  28     """Print an error message to stderr."""
  29     print >> sys.stderr, "ERROR:", msg % args
  30
  31 def warn(msg, *args):
  32     """Print a warning message to stderr."""
  33     print >> sys.stderr, "warning:", msg % args
  34
  35 def die (msg, *args):
  36     """Print as error message to stderr and exit the program."""
  37     error(msg, *args)
  38     sys.exit(1)
  39
  40
  41 class 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
  89 def 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
 106 def 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
 131 def 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
 169 def 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