Merge branch 'gb/am-hg-patch'
[git] / 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