runfiles: Apply repo mapping to Rlocation path (bazel-contrib#998)

fmeum · web-flow · commit 96621391e7e8 · 2023-01-20T06:24:04.000-08:00
diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
@@ -29,7 +29,7 @@ pip.parse(
 )
 use_repo(pip, "pip")
 
-bazel_dep(name = "other_module", version = "")
+bazel_dep(name = "other_module", version = "", repo_name = "our_other_module")
 local_path_override(
     module_name = "other_module",
     path = "other_module",
diff --git a/examples/bzlmod/other_module/other_module/pkg/lib.py b/examples/bzlmod/other_module/other_module/pkg/lib.py
@@ -7,3 +7,7 @@ def GetRunfilePathWithCurrentRepository():
     # For a non-main repository, the name of the runfiles directory is equal to
     # the canonical repository name.
     return r.Rlocation(own_repo + "/other_module/pkg/data/data.txt")
+
+
+def GetRunfilePathWithRepoMapping():
+    return runfiles.Create().Rlocation("other_module/other_module/pkg/data/data.txt")
diff --git a/examples/bzlmod/runfiles/BUILD.bazel b/examples/bzlmod/runfiles/BUILD.bazel
@@ -5,14 +5,14 @@ py_test(
     srcs = ["runfiles_test.py"],
     data = [
         "data/data.txt",
-        "@other_module//other_module/pkg:data/data.txt",
+        "@our_other_module//other_module/pkg:data/data.txt",
     ],
     env = {
         "DATA_RLOCATIONPATH": "$(rlocationpath data/data.txt)",
-        "OTHER_MODULE_DATA_RLOCATIONPATH": "$(rlocationpath @other_module//other_module/pkg:data/data.txt)",
+        "OTHER_MODULE_DATA_RLOCATIONPATH": "$(rlocationpath @our_other_module//other_module/pkg:data/data.txt)",
     },
     deps = [
-        "@other_module//other_module/pkg:lib",
+        "@our_other_module//other_module/pkg:lib",
         "@rules_python//python/runfiles",
     ],
 )
diff --git a/examples/bzlmod/runfiles/runfiles_test.py b/examples/bzlmod/runfiles/runfiles_test.py
@@ -11,12 +11,29 @@ class RunfilesTest(unittest.TestCase):
     def testCurrentRepository(self):
         self.assertEqual(runfiles.Create().CurrentRepository(), "")
 
+    def testRunfilesWithRepoMapping(self):
+        data_path = runfiles.Create().Rlocation("example_bzlmod/runfiles/data/data.txt")
+        with open(data_path) as f:
+            self.assertEqual(f.read().strip(), "Hello, example_bzlmod!")
+
     def testRunfileWithRlocationpath(self):
         data_rlocationpath = os.getenv("DATA_RLOCATIONPATH")
         data_path = runfiles.Create().Rlocation(data_rlocationpath)
         with open(data_path) as f:
             self.assertEqual(f.read().strip(), "Hello, example_bzlmod!")
 
+    def testRunfileInOtherModuleWithOurRepoMapping(self):
+        data_path = runfiles.Create().Rlocation(
+            "our_other_module/other_module/pkg/data/data.txt"
+        )
+        with open(data_path) as f:
+            self.assertEqual(f.read().strip(), "Hello, other_module!")
+
+    def testRunfileInOtherModuleWithItsRepoMapping(self):
+        data_path = lib.GetRunfilePathWithRepoMapping()
+        with open(data_path) as f:
+            self.assertEqual(f.read().strip(), "Hello, other_module!")
+
     def testRunfileInOtherModuleWithCurrentRepository(self):
         data_path = lib.GetRunfilePathWithCurrentRepository()
         with open(data_path) as f:
diff --git a/python/runfiles/runfiles.py b/python/runfiles/runfiles.py
@@ -126,9 +126,12 @@ def __init__(self, strategy):
         # type: (Union[_ManifestBased, _DirectoryBased]) -> None
         self._strategy = strategy
         self._python_runfiles_root = _FindPythonRunfilesRoot()
+        self._repo_mapping = _ParseRepoMapping(
+            strategy.RlocationChecked("_repo_mapping")
+        )
 
-    def Rlocation(self, path):
-        # type: (str) -> Optional[str]
+    def Rlocation(self, path, source_repo=None):
+        # type: (str, Optional[str]) -> Optional[str]
         """Returns the runtime path of a runfile.
 
         Runfiles are data-dependencies of Bazel-built binaries and tests.
@@ -141,6 +144,13 @@ def Rlocation(self, path):
 
         Args:
           path: string; runfiles-root-relative path of the runfile
+          source_repo: string; optional; the canonical name of the repository
+            whose repository mapping should be used to resolve apparent to
+            canonical repository names in `path`. If `None` (default), the
+            repository mapping of the repository containing the caller of this
+            method is used. Explicitly setting this parameter should only be
+            necessary for libraries that want to wrap the runfiles library. Use
+            `CurrentRepository` to obtain canonical repository names.
         Returns:
           the path to the runfile, which the caller should check for existence, or
           None if the method doesn't know about this runfile
@@ -165,7 +175,31 @@ def Rlocation(self, path):
             raise ValueError('path is absolute without a drive letter: "%s"' % path)
         if os.path.isabs(path):
             return path
-        return self._strategy.RlocationChecked(path)
+
+        if source_repo is None and self._repo_mapping:
+            # Look up runfiles using the repository mapping of the caller of the
+            # current method. If the repo mapping is empty, determining this
+            # name is not necessary.
+            source_repo = self.CurrentRepository(frame=2)
+
+        # Split off the first path component, which contains the repository
+        # name (apparent or canonical).
+        target_repo, _, remainder = path.partition("/")
+        if not remainder or (source_repo, target_repo) not in self._repo_mapping:
+            # One of the following is the case:
+            # - not using Bzlmod, so the repository mapping is empty and
+            #   apparent and canonical repository names are the same
+            # - target_repo is already a canonical repository name and does not
+            #   have to be mapped.
+            # - path did not contain a slash and referred to a root symlink,
+            #   which also should not be mapped.
+            return self._strategy.RlocationChecked(path)
+
+        # target_repo is an apparent repository name. Look up the corresponding
+        # canonical repository name with respect to the current repository,
+        # identified by its canonical name.
+        target_canonical = self._repo_mapping[(source_repo, target_repo)]
+        return self._strategy.RlocationChecked(target_canonical + "/" + remainder)
 
     def EnvVars(self):
         # type: () -> Dict[str, str]
@@ -254,6 +288,31 @@ def _FindPythonRunfilesRoot():
     return root
 
 
+def _ParseRepoMapping(repo_mapping_path):
+    # type: (Optional[str]) -> Dict[Tuple[str, str], str]
+    """Parses the repository mapping manifest."""
+    # If the repository mapping file can't be found, that is not an error: We
+    # might be running without Bzlmod enabled or there may not be any runfiles.
+    # In this case, just apply an empty repo mapping.
+    if not repo_mapping_path:
+        return {}
+    try:
+        with open(repo_mapping_path, "r") as f:
+            content = f.read()
+    except FileNotFoundError:
+        return {}
+
+    repo_mapping = {}
+    for line in content.split("\n"):
+        if not line:
+            # Empty line following the last line break
+            break
+        current_canonical, target_local, target_canonical = line.split(",")
+        repo_mapping[(current_canonical, target_local)] = target_canonical
+
+    return repo_mapping
+
+
 class _ManifestBased(object):
     """`Runfiles` strategy that parses a runfiles-manifest to look up runfiles."""
 
diff --git a/tests/runfiles/runfiles_test.py b/tests/runfiles/runfiles_test.py

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ pip.parse(`
`29`	`29`	`)`
`30`	`30`	`use_repo(pip, "pip")`
`31`	`31`
`32`		`-bazel_dep(name = "other_module", version = "")`
	`32`	`+bazel_dep(name = "other_module", version = "", repo_name = "our_other_module")`
`33`	`33`	`local_path_override(`
`34`	`34`	`module_name = "other_module",`
`35`	`35`	`path = "other_module",`