]>
Commit | Line | Data |
---|---|---|
2fe40b63 SR |
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 |