update_checkout: re-read config for cross-repo testing

When updating versions of some dependencies, it's important to be able to re-read the checkout config to get the new versions to actually check them out.

Author: MaxDesiatov

utils/update_checkout/update_checkout/update_checkout.py

@@ -97,6 +97,21 @@ def find_rev_by_timestamp(timestamp, repo_name, refspec):
 
 def get_branch_for_repo(config, repo_name, scheme_name, scheme_map,
                         cross_repos_pr):
+    """Infer, fetch, and return a branch corresponding to a given PR, otherwise
+    return a branch found in the config for this repository name.
+
+    Args:
+        config (Dict[str, Any]): deserialized `update-checkout-config.json`
+        repo_name (str): name of the repository for checking out the branch
+        scheme_name (str): name of the scheme to look up in the config
+        scheme_map (Dict[str, str]): map of repo names to branches to check out
+        cross_repos_pr (Dict[str, str]): map of repo ids to PRs to check out
+
+    Returns:
+        Tuple[str, bool]: a pair of a checked out branch and a boolean
+        indicating whether this repo matched any `cross_repos_pr`.
+    """
+
     cross_repo = False
     repo_branch = scheme_name
     if scheme_map:
@@ -249,17 +264,30 @@ def get_timestamp_to_match(args):
                              echo=False).strip()
 
 
-def update_all_repositories(args, config, scheme_name, cross_repos_pr):
-    scheme_map = None
+def get_scheme_map(config, scheme_name):
+    """Find a mapping from repository IDs to branches in the config.
+
+    Args:
+        config (Dict[str, Any]): deserialized `update-checkout-config.json`
+        scheme_name (str): name of the scheme to look up in `config`
+
+    Returns:
+        Dict[str, str]: a mapping from repos to branches for the given scheme.
+    """
+
     if scheme_name:
         # This loop is only correct, since we know that each alias set has
         # unique contents. This is checked by validate_config. Thus the first
         # branch scheme data that has scheme_name as one of its aliases is
         # the only possible correct answer.
         for v in config['branch-schemes'].values():
             if scheme_name in v['aliases']:
-                scheme_map = v['repos']
-                break
+                return v['repos']
+
+    return None
+
+
+def update_all_repositories(args, config, scheme_name, scheme_map, cross_repos_pr):
     pool_args = []
     timestamp = get_timestamp_to_match(args)
     for repo_name in config['repos'].keys():
@@ -598,22 +626,14 @@ def main():
     clone_with_ssh = args.clone_with_ssh
     skip_history = args.skip_history
     skip_tags = args.skip_tags
-    scheme = args.scheme
+    scheme_name = args.scheme
     github_comment = args.github_comment
     all_repos = args.all_repositories
 
     with open(args.config) as f:
         config = json.load(f)
     validate_config(config)
 
-    if args.dump_hashes:
-        dump_repo_hashes(args, config)
-        return (None, None)
-
-    if args.dump_hashes_config:
-        dump_repo_hashes(args, config, args.dump_hashes_config)
-        return (None, None)
-
     cross_repos_pr = {}
     if github_comment:
         regex_pr = r'(apple/[-a-zA-Z0-9_]+/pull/\d+|apple/[-a-zA-Z0-9_]+#\d+)'
@@ -622,18 +642,46 @@ def main():
         repos_with_pr = [pr.replace('/pull/', '#') for pr in repos_with_pr]
         cross_repos_pr = dict(pr.split('#') for pr in repos_with_pr)
 
+    # If branch is None, default to using the default branch alias
+    # specified by our configuration file.
+    if scheme_name is None:
+        scheme_name = config['default-branch-scheme']
+
+    scheme_map = get_scheme_map(config, scheme_name)
+
+    swift_repo_path = os.path.join(args.source_root, 'swift')
+    if os.path.exists(swift_repo_path):
+        with shell.pushd(swift_repo_path, dry_run=False, echo=True):
+            # Check if `swift` repo itself needs to switch to a cross-repo branch.
+            branch_name, cross_repo = get_branch_for_repo(config, 'swift', 
+                                                          scheme_name,
+                                                          scheme_map,
+                                                          cross_repos_pr)
+
+            if cross_repo:
+                shell.run(['git', 'checkout', branch_name], echo=True,
+                        prefix="[swift] ")
+
+                # Re-read the config after checkout.
+                with open(args.config) as f:
+                    config = json.load(f)
+                validate_config(config)
+
+    if args.dump_hashes:
+        dump_repo_hashes(args, config)
+        return (None, None)
+
+    if args.dump_hashes_config:
+        dump_repo_hashes(args, config, args.dump_hashes_config)
+        return (None, None)
+
     clone_results = None
     if clone or clone_with_ssh:
-        # If branch is None, default to using the default branch alias
-        # specified by our configuration file.
-        if scheme is None:
-            scheme = config['default-branch-scheme']
-
         skip_repo_list = skip_list_for_platform(config, all_repos)
         skip_repo_list.extend(args.skip_repository_list)
         clone_results = obtain_all_additional_swift_sources(args, config,
                                                             clone_with_ssh,
-                                                            scheme,
+                                                            scheme_name,
                                                             skip_history,
                                                             skip_tags,
                                                             skip_repo_list)
@@ -646,8 +694,8 @@ def main():
         print("You don't have all swift sources. "
               "Call this script with --clone to get them.")
 
-    update_results = update_all_repositories(args, config, scheme,
-                                             cross_repos_pr)
+    update_results = update_all_repositories(args, config, scheme_name,
+                                             scheme_map, cross_repos_pr)
     fail_count = 0
     fail_count += check_parallel_results(clone_results, "CLONE")
     fail_count += check_parallel_results(update_results, "UPDATE")