kleaf: bazel.py: use re-written bazelrc for Kleaf sub repo.

When kleaf is used as a submodule of the main repository,
it is referred to as @kleaf. All labels starting with
//build in Kleaf-specific bazelrc files are replaced with
@kleaf//build.

Generated bazelrc files are removed AFTER `bazel clean` is
executed. When one executes `tools/bazel clean`, the wrapper
script executes `bazel clean`, then deletes the generated
bazelrc files.

Test: TH
Bug: 293260835
Change-Id: I546aedf43d74d8ed152566e99851524d2ace809e

1 files changed, 104 insertions, 32 deletions

diff --git a/kleaf/bazel.py b/kleaf/bazel.py
index 06025c3..a65d3b9 100755
--- a/kleaf/bazel.py
+++ b/kleaf/bazel.py
@@ -307,47 +307,48 @@ class BazelWrapper(KleafHelpPrinter):
         self.gen_bazelrc_dir = self.absolute_out_dir / "bazel/bazelrc"
         os.makedirs(self.gen_bazelrc_dir, exist_ok=True)
 
-        self.transformed_startup_options += [
+        self.transformed_startup_options += self._transform_bazelrc_files([
             # Add support for various configs
             # Do not sort, the order here might matter.
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/ants.bazelrc'}",
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/android_ci.bazelrc'}",
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/local.bazelrc'}",
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/fast.bazelrc'}",
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/rbe.bazelrc'}",
-        ]
-
-        self.transformed_startup_options.append(
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/stamp.bazelrc'}",
-        )
-
-        self.transformed_startup_options += [
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/release.bazelrc'}",
-            f"--bazelrc={self.kleaf_repo_dir / FLAGS_BAZEL_RC}",
-        ]
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/ants.bazelrc",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/android_ci.bazelrc",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/local.bazelrc",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/fast.bazelrc",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/rbe.bazelrc",
+        ])
+
+        self.transformed_startup_options += self._transform_bazelrc_files([
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/stamp.bazelrc",
+        ])
+
+        self.transformed_startup_options += self._transform_bazelrc_files([
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/release.bazelrc",
+            self.kleaf_repo_dir / FLAGS_BAZEL_RC,
+        ])
 
         cache_dir_bazelrc = self.gen_bazelrc_dir / "cache_dir.bazelrc"
-        os.makedirs(os.path.dirname(cache_dir_bazelrc), exist_ok=True)
         with open(cache_dir_bazelrc, "w") as f:
+            # The label //build/... will be re-written by _transform_bazelrc_files.
             f.write(textwrap.dedent(f"""\
                 build --//build/kernel/kleaf:cache_dir={shlex.quote(str(self.known_args.cache_dir))}
             """))
 
         if not self.known_startup_options.help:
-            self.transformed_startup_options.append(
-                f"--bazelrc={cache_dir_bazelrc}")
+            self.transformed_startup_options += self._transform_bazelrc_files([
+                cache_dir_bazelrc,
+            ])
 
-        self.transformed_startup_options += [
+        self.transformed_startup_options += self._transform_bazelrc_files([
             # Toolchains and platforms
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/hermetic_cc.bazelrc'}",
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/platforms.bazelrc'}",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/hermetic_cc.bazelrc",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/platforms.bazelrc",
             # Control Network access - with no internet by default.
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/network.bazelrc'}",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/network.bazelrc",
             # Experimental bzlmod support
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/bazelrc/bzlmod.bazelrc'}",
+            self.kleaf_repo_dir / "build/kernel/kleaf/bazelrc/bzlmod.bazelrc",
 
-            f"--bazelrc={self.kleaf_repo_dir / 'build/kernel/kleaf/common.bazelrc'}",
-        ]
+            self.kleaf_repo_dir / "build/kernel/kleaf/common.bazelrc",
+        ])
 
     def _build_final_args(self) -> list[str]:
         """Builds the final arguments for the subprocess."""
@@ -372,6 +373,46 @@ class BazelWrapper(KleafHelpPrinter):
 
         return final_args
 
+    def _transform_bazelrc_files(self, bazelrc_files: list[pathlib.Path]) -> list[str]:
+        """Given a list of bazelrc files, return startup options."""
+        startup_options = []
+        for old_path in bazelrc_files:
+            new_path = self._rewrite_bazelrc_file(old_path)
+            startup_options.append(f"--bazelrc={new_path}")
+        return startup_options
+
+    def _rewrite_bazelrc_file(self, old_path: pathlib.Path) -> pathlib.Path:
+        """Given a bazelrc file, rewrite and return the path."""
+        if self._kleaf_repository_is_top_workspace():
+            # common case; Kleaf tooling is in main Bazel workspace
+            return old_path
+        with open(old_path) as old_file:
+            content = old_file.read()
+
+        # Rewrite //build to @<kleaf_repo_name>//build
+        content = content.replace(
+            "//build", f"{self._kleaf_repo_name()}//build")
+
+        new_path = self.gen_bazelrc_dir / old_path.name
+        os.makedirs(new_path.parent, exist_ok=True)
+        with open(new_path, "w") as new_file:
+            new_file.write(content)
+        return new_path
+
+    def _kleaf_repository_is_top_workspace(self):
+        """Returns true if the Kleaf repository is the top-level workspace @."""
+        return self.workspace_dir == self.kleaf_repo_dir
+
+    def _kleaf_repo_name(self):
+        """Returns the name to the Kleaf repository."""
+        if self._kleaf_repository_is_top_workspace():
+            return "@"
+        # The main repository must refer to the Kleaf repository as @kleaf.
+        # TODO(b/276493276): Once we completely migrate to bzlmod, labels
+        # in bazelrc may be referred to as @kleaf//, then _rewrite_bazelrc_file
+        # may be deleted.
+        return f"@kleaf"
+
     def _print_help(self):
         print("===============================")
 
@@ -433,30 +474,59 @@ class BazelWrapper(KleafHelpPrinter):
         if self.known_startup_options.help or self.command == "help":
             self._print_help()
 
+        # Whether to run bazel comamnd as subprocess
+        run_as_subprocess = False
+        # Regex to filter output / stderr lines
+        filter_regex = None
+        # Epilog coroutine after bazel command finishes
+        epilog_coro = None
+
         if self.known_args.strip_execroot:
-            import asyncio
-            import re
+            run_as_subprocess = True
             if self.absolute_user_root.is_relative_to(self.absolute_out_dir):
                 filter_regex = re.compile(
                     str(self.absolute_out_dir) + r"/\S+?/execroot/__main__/")
             else:
                 filter_regex = re.compile(
                     str(self.absolute_user_root) + r"/\S+?/execroot/__main__/")
-            asyncio.run(run(final_args, self.env, filter_regex))
+
+        if self.command == "clean":
+            run_as_subprocess = True
+            epilog_coro = self.remove_gen_bazelrc_dir()
+
+        if run_as_subprocess:
+            import asyncio
+            import re
+            asyncio.run(run(final_args, self.env, filter_regex, epilog_coro))
         else:
             os.execve(path=self.bazel_path, argv=final_args, env=self.env)
 
+    async def remove_gen_bazelrc_dir(self):
+        sys.stderr.write("INFO: Deleting generated bazelrc directory.\n")
+        shutil.rmtree(self.gen_bazelrc_dir, ignore_errors=True)
+
 
 async def output_filter(input_stream, output_stream, filter_regex):
+    """Pipes input to output, optionally filtering lines with given filter_regex.
+
+    If filter_regex is None, don't filter lines.
+    """
     import re
     while not input_stream.at_eof():
         output = await input_stream.readline()
-        output = re.sub(filter_regex, "", output.decode())
-        output_stream.buffer.write(output.encode())
+        if filter_regex:
+            output = re.sub(filter_regex, "", output.decode()).encode()
+        output_stream.buffer.write(output)
         output_stream.flush()
 
 
-async def run(command, env, filter_regex):
+async def run(command, env, filter_regex, epilog_coro):
+    """Runs command with env asynchronously.
+
+    Outputs are filtered with filter_regex if it is not None.
+
+    At the end, run the coroutine epilog_coro if it is not None.
+    """
     import asyncio
     process = await asyncio.create_subprocess_exec(
         *command,
@@ -470,6 +540,8 @@ async def run(command, env, filter_regex):
         output_filter(process.stdout, sys.stdout, filter_regex),
     )
     await process.wait()
+    if epilog_coro:
+        await epilog_coro
 
 
 if __name__ == "__main__":