diff options
-rw-r--r-- | subcmds/checkout.py | 37 |
1 files changed, 27 insertions, 10 deletions
diff --git a/subcmds/checkout.py b/subcmds/checkout.py index fbb13653..cf54ced7 100644 --- a/subcmds/checkout.py +++ b/subcmds/checkout.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 functools | ||
16 | import multiprocessing | ||
15 | import sys | 17 | import sys |
16 | from command import Command | 18 | |
19 | from command import Command, DEFAULT_LOCAL_JOBS, WORKER_BATCH_SIZE | ||
17 | from progress import Progress | 20 | from progress import Progress |
18 | 21 | ||
19 | 22 | ||
@@ -31,27 +34,41 @@ The command is equivalent to: | |||
31 | 34 | ||
32 | repo forall [<project>...] -c git checkout <branchname> | 35 | repo forall [<project>...] -c git checkout <branchname> |
33 | """ | 36 | """ |
37 | PARALLEL_JOBS = DEFAULT_LOCAL_JOBS | ||
34 | 38 | ||
35 | def ValidateOptions(self, opt, args): | 39 | def ValidateOptions(self, opt, args): |
36 | if not args: | 40 | if not args: |
37 | self.Usage() | 41 | self.Usage() |
38 | 42 | ||
43 | def _ExecuteOne(self, nb, project): | ||
44 | """Checkout one project.""" | ||
45 | return (project.CheckoutBranch(nb), project) | ||
46 | |||
39 | def Execute(self, opt, args): | 47 | def Execute(self, opt, args): |
40 | nb = args[0] | 48 | nb = args[0] |
41 | err = [] | 49 | err = [] |
42 | success = [] | 50 | success = [] |
43 | all_projects = self.GetProjects(args[1:]) | 51 | all_projects = self.GetProjects(args[1:]) |
44 | 52 | ||
45 | pm = Progress('Checkout %s' % nb, len(all_projects)) | 53 | def _ProcessResults(results): |
46 | for project in all_projects: | 54 | for status, project in results: |
47 | pm.update() | 55 | if status is not None: |
56 | if status: | ||
57 | success.append(project) | ||
58 | else: | ||
59 | err.append(project) | ||
60 | pm.update() | ||
48 | 61 | ||
49 | status = project.CheckoutBranch(nb) | 62 | pm = Progress('Checkout %s' % nb, len(all_projects)) |
50 | if status is not None: | 63 | # NB: Multiprocessing is heavy, so don't spin it up for one job. |
51 | if status: | 64 | if len(all_projects) == 1 or opt.jobs == 1: |
52 | success.append(project) | 65 | _ProcessResults(self._ExecuteOne(nb, x) for x in all_projects) |
53 | else: | 66 | else: |
54 | err.append(project) | 67 | with multiprocessing.Pool(opt.jobs) as pool: |
68 | results = pool.imap_unordered( | ||
69 | functools.partial(self._ExecuteOne, nb), all_projects, | ||
70 | chunksize=WORKER_BATCH_SIZE) | ||
71 | _ProcessResults(results) | ||
55 | pm.end() | 72 | pm.end() |
56 | 73 | ||
57 | if err: | 74 | if err: |