Sverre Rabbelier | 2fe40b6 | 2009-11-18 02:42:32 +0100 | [diff] [blame] | 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 |