Add CurrentRepository() to Python runfiles library

`runfiles.CurrentRepository()` can be used to get the canonical name of
the Bazel repository containing the caller at runtime. This information
is required to look up runfiles while taking repository mappings into
account.

src/test/py/bazel/py_test.py

@@ -207,6 +207,108 @@ def testPyTestWithStdlibCollisionRunsRemotely(self):
     self.AssertExitCode(exit_code, 0, stderr, stdout)
     self.assertIn('Test ran', stdout)
 
+class PyRunfilesLibraryTest(test_base.TestBase):
+  def testPyRunfilesLibraryCurrentRepository(self):
+    self.CreateWorkspaceWithDefaultRepos('WORKSPACE', [
+        'local_repository(',
+        '  name = "other_repo",',
+        '  path = "other_repo_path",',
+        ')'
+    ])
+
+    self.ScratchFile('pkg/BUILD.bazel', [
+        'py_library(',
+        '  name = "library",',
+        '  srcs = ["library.py"],',
+        '  visibility = ["//visibility:public"],',
+        '  deps = ["@bazel_tools//tools/python/runfiles"],',
+        ')',
+        '',
+        'py_binary(',
+        '  name = "binary",',
+        '  srcs = ["binary.py"],',
+        '  deps = [',
+        '    ":library",',
+        '    "@bazel_tools//tools/python/runfiles",'
+        '  ],',
+        ')',
+        '',
+        'py_test(',
+        '  name = "test",',
+        '  srcs = ["test.py"],',
+        '  deps = [',
+        '    ":library",',
+        '    "@bazel_tools//tools/python/runfiles",',
+        '  ],',
+        ')',
+    ])
+    self.ScratchFile('pkg/library.py', [
+        'from bazel_tools.tools.python.runfiles import runfiles',
+        'def print_repo_name():',
+        '  print("in pkg/library.py: \'%s\'" % runfiles.CurrentRepository())',
+    ])
+    self.ScratchFile('pkg/binary.py', [
+        'from bazel_tools.tools.python.runfiles import runfiles',
+        'from pkg import library',
+        'library.print_repo_name()',
+        'print("in pkg/binary.py: \'%s\'" % runfiles.CurrentRepository())',
+    ])
+    self.ScratchFile('pkg/test.py', [
+        'from bazel_tools.tools.python.runfiles import runfiles',
+        'from pkg import library',
+        'library.print_repo_name()',
+        'print("in pkg/test.py: \'%s\'" % runfiles.CurrentRepository())',
+    ])
+
+    self.ScratchFile('other_repo_path/WORKSPACE')
+    self.ScratchFile('other_repo_path/pkg/BUILD.bazel', [
+        'py_binary(',
+        '  name = "binary",',
+        '  srcs = ["binary.py"],',
+        '  deps = [',
+        '    "@//pkg:library",',
+        '    "@bazel_tools//tools/python/runfiles",'
+        '  ],',
+        ')',
+        '',
+        'py_test(',
+        '  name = "test",',
+        '  srcs = ["test.py"],',
+        '  deps = [',
+        '    "@//pkg:library",',
+        '    "@bazel_tools//tools/python/runfiles",',
+        '  ],',
+        ')',
+    ])
+    self.ScratchFile('other_repo_path/pkg/binary.py', [
+        'from bazel_tools.tools.python.runfiles import runfiles',
+        'from pkg import library',
+        'library.print_repo_name()',
+        'print("in external/other_repo/pkg/binary.py: \'%s\'" % runfiles.CurrentRepository())',
+    ])
+    self.ScratchFile('other_repo_path/pkg/test.py', [
+        'from bazel_tools.tools.python.runfiles import runfiles',
+        'from pkg import library',
+        'library.print_repo_name()',
+        'print("in external/other_repo/pkg/test.py: \'%s\'" % runfiles.CurrentRepository())',
+    ])
+
+    _, stdout, _ = self.RunBazel(['run', '//pkg:binary'])
+    self.assertIn('in pkg/binary.py: \'\'', stdout)
+    self.assertIn('in pkg/library.py: \'\'', stdout)
+
+    _, stdout, _ = self.RunBazel(['test', '//pkg:test', '--test_output=streamed'])
+    self.assertIn('in pkg/test.py: \'\'', stdout)
+    self.assertIn('in pkg/library.py: \'\'', stdout)
+
+    _, stdout, _ = self.RunBazel(['run', '@other_repo//pkg:binary'])
+    self.assertIn('in external/other_repo/pkg/binary.py: \'other_repo\'', stdout)
+    self.assertIn('in pkg/library.py: \'\'', stdout)
+
+    _, stdout, _ = self.RunBazel(['test', '@other_repo//pkg:test', '--test_output=streamed'])
+    self.assertIn('in external/other_repo/pkg/test.py: \'other_repo\'', stdout)
+    self.assertIn('in pkg/library.py: \'\'', stdout)
+
 
 if __name__ == '__main__':
   unittest.main()

tools/python/gen_runfiles_constants.bzl

@@ -0,0 +1,30 @@
+# Copyright 2022 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+_RUNFILES_CONSTANTS_TEMPLATE = """# The name of the runfiles directory corresponding to the main repository.
+MAIN_REPOSITORY_RUNFILES_DIRECTORY = '%s'
+"""
+
+def _gen_runfiles_constants_impl(ctx):
+    out = ctx.actions.declare_file(ctx.attr.name + ".py")
+    ctx.actions.write(out, _RUNFILES_CONSTANTS_TEMPLATE % ctx.workspace_name)
+
+    return DefaultInfo(
+        files = depset([out]),
+        runfiles = ctx.runfiles([out]),
+    )
+
+gen_runfiles_constants = rule(
+    implementation = _gen_runfiles_constants_impl,
+)

tools/python/runfiles/BUILD

@@ -1,3 +1,4 @@
+load("//tools/python:gen_runfiles_constants.bzl", "gen_runfiles_constants")
 load("//tools/python:private/defs.bzl", "py_library", "py_test")
 
 package(default_visibility = ["//visibility:private"])
@@ -20,7 +21,14 @@ filegroup(
 py_library(
     name = "runfiles",
     testonly = 1,
-    srcs = ["runfiles.py"],
+    srcs = [
+        "runfiles.py",
+        ":runfiles_constants",
+    ],
+)
+
+gen_runfiles_constants(
+    name = "runfiles_constants",
 )
 
 py_test(

tools/python/runfiles/BUILD.tools

@@ -1,7 +1,15 @@
+load("//tools/python:gen_runfiles_constants.bzl", "gen_runfiles_constants")
 load("//tools/python:private/defs.bzl", "py_library")
 
 py_library(
     name = "runfiles",
-    srcs = ["runfiles.py"],
+    srcs = [
+        "runfiles.py",
+        ":runfiles_constants",
+    ],
     visibility = ["//visibility:public"],
 )
+
+gen_runfiles_constants(
+    name = "runfiles_constants",
+)

tools/python/runfiles/runfiles.py

@@ -58,9 +58,12 @@
       p = subprocess.Popen([r.Rlocation("path/to/binary")], env, ...)
 """
 
+import inspect
 import os
 import posixpath
+import sys
 
+from .runfiles_constants import MAIN_REPOSITORY_RUNFILES_DIRECTORY
 
 def CreateManifestBased(manifest_path):
   return _Runfiles(_ManifestBased(manifest_path))
@@ -114,6 +117,7 @@ class _Runfiles(object):
 
   def __init__(self, strategy):
     self._strategy = strategy
+    self._python_runfiles_root = _FindPythonRunfilesRoot()
 
   def Rlocation(self, path):
     """Returns the runtime path of a runfile.
@@ -161,6 +165,51 @@ def EnvVars(self):
     """
     return self._strategy.EnvVars()
 
+  def CurrentRepository(self, frame = 1):
+    """Returns the canonical name of the caller's Bazel repository.
+
+    This function uses the 'inspect' package to obtain the file containing the
+    frame-th caller of this method and returns the canonical name of the Bazel
+    repository containing that file.
+
+    For example, this function returns '' (the empty string) when called from
+    the main repository and 'bazel_tools' when called from the @bazel_tools
+    repository.
+
+    Args:
+      frame: int; the stack frame to return the repository name for. Defaults to
+       1, the caller of the CurrentRepository function.
+    Returns:
+      the canonical name of the Bazel repository containing the file containing
+      the frame-th caller of this function
+    Raises:
+      ValueError: if the caller cannot be determined or the caller's file path
+       is not contained in the Python runfiles tree
+    """
+    caller_path = inspect.getfile(sys._getframe(frame))
+    caller_runfiles_path = os.path.relpath(caller_path, self._python_runfiles_root)
+    if caller_runfiles_path.startswith(".." + os.path.sep):
+      raise ValueError('{} does not lie under the runfiles root {}'.format(caller_path, self._python_runfiles_root))
+
+    caller_runfiles_directory = caller_runfiles_path[:caller_runfiles_path.find(os.path.sep)]
+    if caller_runfiles_directory == MAIN_REPOSITORY_RUNFILES_DIRECTORY:
+      # The canonical name of the main repository (also know as the workspace)
+      # is the empty string.
+      return ''
+    # For all other repositories, the name of the runfiles directory is the
+    # canonical name.
+    return caller_runfiles_directory
+
+def _FindPythonRunfilesRoot():
+  root = __file__
+  # Walk up our own runfiles path to the root of the runfiles tree from which
+  # the current file is being run. This path coincides with what the Bazel
+  # Python stub sets up as sys.path[0]. Since that entry can be changed at
+  # runtime, we rederive it here.
+  for _ in range("bazel_tools/tools/python/runfiles/runfiles.py".count("/") + 1):
+    root = os.path.dirname(root)
+  return root
+
 
 class _ManifestBased(object):
   """`Runfiles` strategy that parses a runfiles-manifest to look up runfiles."""

tools/python/runfiles/runfiles_test.py

@@ -262,6 +262,9 @@ def testPathsFromEnvvars(self):
     self.assertEqual(mf, "")
     self.assertEqual(dr, "")
 
+  def testCurrentRepository(self):
+    self.assertEqual(runfiles.Create().CurrentRepository(), "")
+
   @staticmethod
   def IsWindows():
     return os.name == "nt"