diff options
Diffstat (limited to 'subcmds/prune.py')
-rw-r--r-- | subcmds/prune.py | 20 |
1 files changed, 8 insertions, 12 deletions
diff --git a/subcmds/prune.py b/subcmds/prune.py index 4084c8b6..236b647f 100644 --- a/subcmds/prune.py +++ b/subcmds/prune.py | |||
@@ -13,10 +13,9 @@ | |||
13 | # limitations under the License. | 13 | # limitations under the License. |
14 | 14 | ||
15 | import itertools | 15 | import itertools |
16 | import multiprocessing | ||
17 | 16 | ||
18 | from color import Coloring | 17 | from color import Coloring |
19 | from command import DEFAULT_LOCAL_JOBS, PagedCommand, WORKER_BATCH_SIZE | 18 | from command import DEFAULT_LOCAL_JOBS, PagedCommand |
20 | 19 | ||
21 | 20 | ||
22 | class Prune(PagedCommand): | 21 | class Prune(PagedCommand): |
@@ -36,18 +35,15 @@ class Prune(PagedCommand): | |||
36 | 35 | ||
37 | # NB: Should be able to refactor this module to display summary as results | 36 | # NB: Should be able to refactor this module to display summary as results |
38 | # come back from children. | 37 | # come back from children. |
39 | def _ProcessResults(results): | 38 | def _ProcessResults(_pool, _output, results): |
40 | return list(itertools.chain.from_iterable(results)) | 39 | return list(itertools.chain.from_iterable(results)) |
41 | 40 | ||
42 | # NB: Multiprocessing is heavy, so don't spin it up for one job. | 41 | all_branches = self.ExecuteInParallel( |
43 | if len(projects) == 1 or opt.jobs == 1: | 42 | opt.jobs, |
44 | all_branches = _ProcessResults(self._ExecuteOne(x) for x in projects) | 43 | self._ExecuteOne, |
45 | else: | 44 | projects, |
46 | with multiprocessing.Pool(opt.jobs) as pool: | 45 | callback=_ProcessResults, |
47 | results = pool.imap( | 46 | ordered=True) |
48 | self._ExecuteOne, projects, | ||
49 | chunksize=WORKER_BATCH_SIZE) | ||
50 | all_branches = _ProcessResults(results) | ||
51 | 47 | ||
52 | if not all_branches: | 48 | if not all_branches: |
53 | return | 49 | return |