diff options
Diffstat (limited to 'subcmds')
| -rw-r--r-- | subcmds/prune.py | 29 |
1 files changed, 25 insertions, 4 deletions
diff --git a/subcmds/prune.py b/subcmds/prune.py index 8cad8122..4084c8b6 100644 --- a/subcmds/prune.py +++ b/subcmds/prune.py | |||
| @@ -12,8 +12,11 @@ | |||
| 12 | # See the License for the specific language governing permissions and | 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | 14 | ||
| 15 | import itertools | ||
| 16 | import multiprocessing | ||
| 17 | |||
| 15 | from color import Coloring | 18 | from color import Coloring |
| 16 | from command import PagedCommand | 19 | from command import DEFAULT_LOCAL_JOBS, PagedCommand, WORKER_BATCH_SIZE |
| 17 | 20 | ||
| 18 | 21 | ||
| 19 | class Prune(PagedCommand): | 22 | class Prune(PagedCommand): |
| @@ -22,11 +25,29 @@ class Prune(PagedCommand): | |||
| 22 | helpUsage = """ | 25 | helpUsage = """ |
| 23 | %prog [<project>...] | 26 | %prog [<project>...] |
| 24 | """ | 27 | """ |
| 28 | PARALLEL_JOBS = DEFAULT_LOCAL_JOBS | ||
| 29 | |||
| 30 | def _ExecuteOne(self, project): | ||
| 31 | """Process one project.""" | ||
| 32 | return project.PruneHeads() | ||
| 25 | 33 | ||
| 26 | def Execute(self, opt, args): | 34 | def Execute(self, opt, args): |
| 27 | all_branches = [] | 35 | projects = self.GetProjects(args) |
| 28 | for project in self.GetProjects(args): | 36 | |
| 29 | all_branches.extend(project.PruneHeads()) | 37 | # NB: Should be able to refactor this module to display summary as results |
| 38 | # come back from children. | ||
| 39 | def _ProcessResults(results): | ||
| 40 | return list(itertools.chain.from_iterable(results)) | ||
| 41 | |||
| 42 | # NB: Multiprocessing is heavy, so don't spin it up for one job. | ||
| 43 | if len(projects) == 1 or opt.jobs == 1: | ||
| 44 | all_branches = _ProcessResults(self._ExecuteOne(x) for x in projects) | ||
| 45 | else: | ||
| 46 | with multiprocessing.Pool(opt.jobs) as pool: | ||
| 47 | results = pool.imap( | ||
| 48 | self._ExecuteOne, projects, | ||
| 49 | chunksize=WORKER_BATCH_SIZE) | ||
| 50 | all_branches = _ProcessResults(results) | ||
| 30 | 51 | ||
| 31 | if not all_branches: | 52 | if not all_branches: |
| 32 | return | 53 | return |
