OLD | NEW |
---|---|
(Empty) | |
1 # Copyright (c) 2013 The Chromium Authors. All rights reserved. | |
2 # Use of this source code is governed by a BSD-style license that can be | |
3 # found in the LICENSE file. | |
4 | |
5 # Monkeypatch IMapIterator so that Ctrl-C can kill everything properly. | |
6 # Derived from https://gist.github.com/aljungberg/626518 | |
7 import multiprocessing.pool | |
8 from multiprocessing.pool import IMapIterator | |
9 def wrapper(func): | |
10 def wrap(self, timeout=None): | |
11 return func(self, timeout=timeout or 1e100) | |
12 return wrap | |
13 IMapIterator.next = wrapper(IMapIterator.next) | |
14 IMapIterator.__next__ = IMapIterator.next | |
15 | |
16 | |
17 import binascii | |
18 import contextlib | |
19 import functools | |
20 import logging | |
21 import signal | |
22 import subprocess | |
23 import sys | |
24 import tempfile | |
25 import threading | |
26 | |
27 | |
28 # This converts a binary git hash prefix into a posix path, one folder per byte | |
29 # e.g. "\xDE\xAD" -> "de/ad" | |
30 pathlify = lambda s: '/'.join('%02x' % ord(b) for b in s) | |
M-A Ruel
2013/11/07 20:59:11
This seems to create very deep directory. What is
iannucci
2013/11/07 21:44:57
Like the comment above says, it' converts a prefix
| |
31 | |
32 | |
33 GIT_EXE = 'git.bat' if sys.platform.startswith('win') else 'git' | |
34 | |
35 | |
36 class CalledProcessError(Exception): | |
37 def __init__(self, returncode, cmd): | |
38 super(CalledProcessError, self).__init__() | |
39 self.returncode = returncode | |
40 self.cmd = cmd | |
41 | |
42 def __str__(self): | |
43 return ( | |
44 'Command "%s" returned non-zero exit status %d' % | |
45 (self.cmd, self.returncode)) | |
46 | |
47 | |
48 def memoize_one(f): | |
49 """Memoizes a single-argument pure function. | |
50 | |
51 Values of None are not cached. | |
52 | |
53 Adds a mutable attribute to the decorated function: | |
54 * cache (dict) - Maps arg to f(arg) | |
M-A Ruel
2013/11/07 20:59:11
I'd be good to document how one is expected to dis
iannucci
2013/11/07 21:44:57
Done.
| |
55 """ | |
56 cache = {} | |
57 | |
58 @functools.wraps(f) | |
59 def inner(arg): | |
60 ret = cache.get(arg) | |
61 if ret is None: | |
62 ret = f(arg) | |
63 if ret is not None: | |
64 cache[arg] = ret | |
65 return ret | |
66 inner.cache = cache | |
67 | |
68 return inner | |
69 | |
70 | |
71 def _ScopedPool_initer(orig, orig_args): # pragma: no cover | |
72 """Initializer method for ScopedPool's subprocesses. | |
73 | |
74 This helps ScopedPool handle Ctrl-C's correctly. | |
75 """ | |
76 signal.signal(signal.SIGINT, signal.SIG_IGN) | |
77 if orig: | |
78 orig(*orig_args) | |
79 | |
80 | |
81 @contextlib.contextmanager | |
82 def ScopedPool(*args, **kwargs): | |
83 if kwargs.pop('kind', None) == 'threads': | |
84 pool = multiprocessing.pool.ThreadPool(*args, **kwargs) | |
85 else: | |
86 orig, orig_args = kwargs.get('initializer'), kwargs.get('initargs', ()) | |
87 kwargs['initializer'] = _ScopedPool_initer | |
88 kwargs['initargs'] = orig, orig_args | |
89 pool = multiprocessing.pool.Pool(*args, **kwargs) | |
90 | |
91 try: | |
92 yield pool | |
93 pool.close() | |
94 except: | |
95 pool.terminate() | |
96 raise | |
97 finally: | |
98 pool.join() | |
99 | |
100 | |
101 class ProgressPrinter(object): | |
102 """Threaded single-stat status message printer.""" | |
103 def __init__(self, fmt, enabled=None, stream=sys.stderr, period=0.5): | |
104 """Create a ProgressPrinter. | |
105 | |
106 Use it as a context manager which produces a simple 'increment' method: | |
107 | |
108 with ProgressPrinter('(%%(count)d/%d)' % 1000) as inc: | |
109 for i in xrange(1000): | |
110 # do stuff | |
111 if i % 10 == 0: | |
112 inc(10) | |
113 | |
114 Args: | |
115 fmt - String format with a single '%(count)d' where the counter value | |
116 should go. | |
117 enabled (bool) - If this is None, will default to True if | |
118 logging.getLogger() is set to INFO or more verbose. | |
119 stream (file-like) - The stream to print status messages to. | |
120 period (float) - The time in seconds for the printer thread to wait | |
121 between printing. | |
122 """ | |
123 self.fmt = fmt | |
124 if enabled is None: # pragma: no cover | |
125 self.enabled = logging.getLogger().isEnabledFor(logging.INFO) | |
126 else: | |
127 self.enabled = enabled | |
128 | |
129 self._count = 0 | |
130 self._dead = False | |
131 self._dead_cond = threading.Condition() | |
132 self._stream = stream | |
133 self._thread = threading.Thread(target=self._run) | |
134 self._period = period | |
135 | |
136 def _emit(self, s): | |
137 if self.enabled: | |
138 self._stream.write('\r'+s) | |
139 self._stream.flush() | |
140 | |
141 def _run(self): | |
142 with self._dead_cond: | |
143 while not self._dead: | |
144 self._emit(self.fmt % {'count': self._count}) | |
145 self._dead_cond.wait(self._period) | |
146 self._emit((self.fmt+'\n') % {'count': self._count}) | |
147 | |
148 def inc(self, amount=1): | |
149 self._count += amount | |
150 | |
151 def __enter__(self): | |
152 self._thread.start() | |
153 return self.inc | |
154 | |
155 def __exit__(self, _exc_type, _exc_value, _traceback): | |
156 self._dead = True | |
157 with self._dead_cond: | |
158 self._dead_cond.notifyAll() | |
159 self._thread.join() | |
160 del self._thread | |
161 | |
162 | |
163 def parse_committishes(*committishes): | |
164 """This takes one or more committishes, and returns the binary-encoded git | |
165 hashes for them. | |
166 | |
167 A committish is anything which can resolve to a commit. Popular examples: | |
168 * "HEAD" | |
169 * "origin/master" | |
170 * "cool_branch~2" | |
171 | |
172 etc. | |
173 """ | |
174 try: | |
175 return map(binascii.unhexlify, hashes(*committishes)) | |
176 except CalledProcessError: | |
177 raise Exception('one of %s does not seem to be a valid commitish.' % | |
178 str(committishes)) | |
179 | |
180 | |
181 def _check_output(*popenargs, **kwargs): | |
182 """Run a Popen command, and return the stdout as a string. | |
183 | |
184 Throws CalledProcessError if the command returns non-zero. | |
185 | |
186 kwargs: | |
187 indata (str) - Data to provide to the command on stdin. Mutually exclusive | |
188 with the Popen kwarg 'stdin'. | |
189 | |
190 Other than that, popenargs is *args to Popen, and **kwargs is... **kwargs to | |
191 Popen. | |
192 """ | |
193 kwargs.setdefault('stdout', subprocess.PIPE) | |
194 kwargs.setdefault('stderr', subprocess.PIPE) | |
195 indata = kwargs.pop('indata', None) | |
196 if indata is not None: | |
197 kwargs['stdin'] = subprocess.PIPE | |
198 process = subprocess.Popen(*popenargs, **kwargs) | |
199 output, _ = process.communicate(indata) | |
200 if process.returncode: | |
201 cmd = kwargs.get('args') | |
202 if cmd is None: | |
203 cmd = popenargs[0] | |
204 raise CalledProcessError(process.returncode, cmd) | |
205 return output | |
206 | |
207 | |
208 def run(*cmd, **kwargs): | |
209 """Runs a git command. Returns stdout as a string. | |
210 | |
211 If logging is DEBUG, we'll print the command before we run it. | |
212 | |
213 Output string is always strip()'d. | |
214 """ | |
215 cmd = (GIT_EXE,) + cmd | |
216 logging.debug('running: %s', " ".join(repr(tok) for tok in cmd)) | |
217 ret = _check_output(cmd, **kwargs) | |
218 ret = (ret or '').strip() | |
219 return ret | |
220 | |
221 | |
222 def hashes(*reflike): | |
223 return run('rev-parse', *reflike).splitlines() | |
224 | |
225 | |
226 def intern_f(f, kind='blob'): | |
227 """Interns a file object into the git object store. | |
228 | |
229 Args: | |
230 f (file-like object) - The file-like object to intern | |
231 kind (git object type) - One of 'blob', 'commit', 'tree', 'tag'. | |
232 | |
233 Returns the git hash of the interned object (hex encoded). | |
234 """ | |
235 ret = run('hash-object', '-t', kind, '-w', '--stdin', stdin=f) | |
236 f.close() | |
237 return ret | |
238 | |
239 | |
240 def tree(treeish, recurse=False): | |
241 """ | |
242 Args: | |
243 treeish - a git name which resolves to a tree (or to a commit). | |
244 recurse - include just this tree, or all of its decendants too. | |
245 | |
246 Returns a dict formatted like: | |
247 { 'file_name': (mode, type, ref) } | |
248 | |
249 mode is an integer where: | |
250 * 0040000 - Directory | |
251 * 0100644 - Regular non-executable file | |
252 * 0100664 - Regular non-executable group-writeable file | |
253 * 0100755 - Regular executable file | |
254 * 0120000 - Symbolic link | |
255 * 0160000 - Gitlink | |
256 | |
257 type is a string where it's one of 'blob', 'commit', 'tree', 'tag'. | |
258 | |
259 ref is the hex encoded hash of the entry. | |
260 """ | |
261 ret = {} | |
262 opts = ['ls-tree', '--full-tree'] | |
263 if recurse: | |
264 opts += ['-r'] | |
265 opts.append(treeish) | |
266 try: | |
267 for line in run(*opts).splitlines(): | |
268 mode, typ, ref, name = line.split(None, 3) | |
269 ret[name] = (mode, typ, ref) | |
270 except CalledProcessError: | |
271 return None | |
272 return ret | |
273 | |
274 | |
275 def mktree(treedict): | |
276 """Make a git tree object and return its hash. | |
277 | |
278 See tree for the values of mode, type, and ref. | |
279 | |
280 Args: | |
281 treedict - { name: (mode, type, ref) } | |
282 """ | |
283 with tempfile.TemporaryFile() as f: | |
284 for name, (mode, typ, ref) in treedict.iteritems(): | |
285 f.write('%s %s %s\t%s\0' % (mode, typ, ref, name)) | |
286 f.seek(0) | |
287 return run('mktree', '-z', stdin=f) | |
OLD | NEW |