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 @@ # See the License for the specific language governing permissions and # limitations under the License. +import itertools +import multiprocessing + from color import Coloring -from command import PagedCommand +from command import DEFAULT_LOCAL_JOBS, PagedCommand, WORKER_BATCH_SIZE class Prune(PagedCommand): @@ -22,11 +25,29 @@ class Prune(PagedCommand): helpUsage = """ %prog [...] """ + PARALLEL_JOBS = DEFAULT_LOCAL_JOBS + + def _ExecuteOne(self, project): + """Process one project.""" + return project.PruneHeads() def Execute(self, opt, args): - all_branches = [] - for project in self.GetProjects(args): - all_branches.extend(project.PruneHeads()) + projects = self.GetProjects(args) + + # NB: Should be able to refactor this module to display summary as results + # come back from children. + def _ProcessResults(results): + return list(itertools.chain.from_iterable(results)) + + # NB: Multiprocessing is heavy, so don't spin it up for one job. + if len(projects) == 1 or opt.jobs == 1: + all_branches = _ProcessResults(self._ExecuteOne(x) for x in projects) + else: + with multiprocessing.Pool(opt.jobs) as pool: + results = pool.imap( + self._ExecuteOne, projects, + chunksize=WORKER_BATCH_SIZE) + all_branches = _ProcessResults(results) if not all_branches: return