1 files changed, 51 insertions, 50 deletions

diff --git a/gitc_utils.py b/gitc_utils.py
index b47e181c..486bbeb0 100644
--- a/gitc_utils.py
+++ b/gitc_utils.py
@@ -1,5 +1,3 @@
-# -*- coding:utf-8 -*-
-#
 # Copyright (C) 2015 The Android Open Source Project
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import print_function
 import os
+import multiprocessing
 import platform
 import re
 import sys
@@ -29,12 +27,24 @@ from error import ManifestParseError
 
 NUM_BATCH_RETRIEVE_REVISIONID = 32
 
+
 def get_gitc_manifest_dir():
   return wrapper.Wrapper().get_gitc_manifest_dir()
 
+
 def parse_clientdir(gitc_fs_path):
   return wrapper.Wrapper().gitc_parse_clientdir(gitc_fs_path)
 
+
+def _get_project_revision(args):
+  """Worker for _set_project_revisions to lookup one project remote."""
+  (i, url, expr) = args
+  gitcmd = git_command.GitCommand(
+      None, ['ls-remote', url, expr], capture_stdout=True, cwd='/tmp')
+  rc = gitcmd.Wait()
+  return (i, rc, gitcmd.stdout.split('\t', 1)[0])
+
+
 def _set_project_revisions(projects):
   """Sets the revisionExpr for a list of projects.
 
@@ -42,47 +52,38 @@ def _set_project_revisions(projects):
   should not be overly large. Recommend calling this function multiple times
   with each call not exceeding NUM_BATCH_RETRIEVE_REVISIONID projects.
 
-  @param projects: List of project objects to set the revionExpr for.
+  Args:
+    projects: List of project objects to set the revionExpr for.
   """
   # Retrieve the commit id for each project based off of it's current
   # revisionExpr and it is not already a commit id.
-  project_gitcmds = [(
-      project, git_command.GitCommand(None,
-                                      ['ls-remote',
-                                       project.remote.url,
-                                       project.revisionExpr],
-                                      capture_stdout=True, cwd='/tmp'))
-      for project in projects if not git_config.IsId(project.revisionExpr)]
-  for proj, gitcmd in project_gitcmds:
-    if gitcmd.Wait():
-      print('FATAL: Failed to retrieve revisionExpr for %s' % proj)
-      sys.exit(1)
-    revisionExpr = gitcmd.stdout.split('\t')[0]
-    if not revisionExpr:
-      raise ManifestParseError('Invalid SHA-1 revision project %s (%s)' %
-                               (proj.remote.url, proj.revisionExpr))
-    proj.revisionExpr = revisionExpr
-
-def _manifest_groups(manifest):
-  """Returns the manifest group string that should be synced
-
-  This is the same logic used by Command.GetProjects(), which is used during
-  repo sync
-
-  @param manifest: The XmlManifest object
-  """
-  mp = manifest.manifestProject
-  groups = mp.config.GetString('manifest.groups')
-  if not groups:
-    groups = 'default,platform-' + platform.system().lower()
-  return groups
+  with multiprocessing.Pool(NUM_BATCH_RETRIEVE_REVISIONID) as pool:
+    results_iter = pool.imap_unordered(
+        _get_project_revision,
+        ((i, project.remote.url, project.revisionExpr)
+         for i, project in enumerate(projects)
+         if not git_config.IsId(project.revisionExpr)),
+        chunksize=8)
+    for (i, rc, revisionExpr) in results_iter:
+      project = projects[i]
+      if rc:
+        print('FATAL: Failed to retrieve revisionExpr for %s' % project.name)
+        pool.terminate()
+        sys.exit(1)
+      if not revisionExpr:
+        pool.terminate()
+        raise ManifestParseError('Invalid SHA-1 revision project %s (%s)' %
+                                 (project.remote.url, project.revisionExpr))
+      project.revisionExpr = revisionExpr
+
 
 def generate_gitc_manifest(gitc_manifest, manifest, paths=None):
   """Generate a manifest for shafsd to use for this GITC client.
 
-  @param gitc_manifest: Current gitc manifest, or None if there isn't one yet.
-  @param manifest: A GitcManifest object loaded with the current repo manifest.
-  @param paths: List of project paths we want to update.
+  Args:
+    gitc_manifest: Current gitc manifest, or None if there isn't one yet.
+    manifest: A GitcManifest object loaded with the current repo manifest.
+    paths: List of project paths we want to update.
   """
 
   print('Generating GITC Manifest by fetching revision SHAs for each '
@@ -90,7 +91,7 @@ def generate_gitc_manifest(gitc_manifest, manifest, paths=None):
   if paths is None:
     paths = list(manifest.paths.keys())
 
-  groups = [x for x in re.split(r'[,\s]+', _manifest_groups(manifest)) if x]
+  groups = [x for x in re.split(r'[,\s]+', manifest.GetGroupsStr()) if x]
 
   # Convert the paths to projects, and filter them to the matched groups.
   projects = [manifest.paths[p] for p in paths]
@@ -104,11 +105,11 @@ def generate_gitc_manifest(gitc_manifest, manifest, paths=None):
       if not proj.upstream and not git_config.IsId(proj.revisionExpr):
         proj.upstream = proj.revisionExpr
 
-      if not path in gitc_manifest.paths:
+      if path not in gitc_manifest.paths:
         # Any new projects need their first revision, even if we weren't asked
         # for them.
         projects.append(proj)
-      elif not path in paths:
+      elif path not in paths:
         # And copy revisions from the previous manifest if we're not updating
         # them now.
         gitc_proj = gitc_manifest.paths[path]
@@ -118,11 +119,7 @@ def generate_gitc_manifest(gitc_manifest, manifest, paths=None):
         else:
           proj.revisionExpr = gitc_proj.revisionExpr
 
-  index = 0
-  while index < len(projects):
-    _set_project_revisions(
-        projects[index:(index+NUM_BATCH_RETRIEVE_REVISIONID)])
-    index += NUM_BATCH_RETRIEVE_REVISIONID
+  _set_project_revisions(projects)
 
   if gitc_manifest is not None:
     for path, proj in gitc_manifest.paths.items():
@@ -140,16 +137,20 @@ def generate_gitc_manifest(gitc_manifest, manifest, paths=None):
   # Save the manifest.
   save_manifest(manifest)
 
+
 def save_manifest(manifest, client_dir=None):
   """Save the manifest file in the client_dir.
 
-  @param client_dir: Client directory to save the manifest in.
-  @param manifest: Manifest object to save.
+  Args:
+    manifest: Manifest object to save.
+    client_dir: Client directory to save the manifest in.
   """
   if not client_dir:
-    client_dir = manifest.gitc_client_dir
-  with open(os.path.join(client_dir, '.manifest'), 'w') as f:
-    manifest.Save(f, groups=_manifest_groups(manifest))
+    manifest_file = manifest.manifestFile
+  else:
+    manifest_file = os.path.join(client_dir, '.manifest')
+  with open(manifest_file, 'w') as f:
+    manifest.Save(f, groups=manifest.GetGroupsStr())
   # TODO(sbasi/jorg): Come up with a solution to remove the sleep below.
   # Give the GITC filesystem time to register the manifest changes.
   time.sleep(3)