fix(py_venv): patch bazel-runfiles root detection

e2e/MODULE.bazel

@@ -121,6 +121,7 @@ include("//cases/uv-console-script-extras:setup.MODULE.bazel")
 include("//cases/uv-dep-hashes:setup.MODULE.bazel")
 include("//cases/uv-deps-650:setup.MODULE.bazel")
 include("//cases/uv-data-purelib:setup.MODULE.bazel")
+include("//cases/uv-bazel-runfiles-root:setup.MODULE.bazel")
 include("//cases/uv-dup-lock-records:setup.MODULE.bazel")
 include("//cases/uv-include-group:setup.MODULE.bazel")
 include("//cases/uv-incompatible-wheel-source:setup.MODULE.bazel")

e2e/cases/uv-bazel-runfiles-root/BUILD.bazel

@@ -0,0 +1,10 @@
+load("@aspect_rules_py//py:defs.bzl", "py_test")
+
+py_test(
+    name = "test_runfiles_root",
+    srcs = ["test_runfiles_root.py"],
+    data = ["runfiles_root_data.txt"],
+    dep_group = "bazel-runfiles-root",
+    main = "test_runfiles_root.py",
+    deps = ["@pypi_uv_bazel_runfiles_root//bazel_runfiles"],
+)

e2e/cases/uv-bazel-runfiles-root/pyproject.toml

@@ -0,0 +1,8 @@
+[project]
+name = "bazel-runfiles-root"
+version = "0.0.0"
+authors = []
+requires-python = ">=3.11"
+dependencies = [
+    "bazel-runfiles==1.1.0",
+]

e2e/cases/uv-bazel-runfiles-root/runfiles_root_data.txt

@@ -0,0 +1 @@
+runfiles root data

e2e/cases/uv-bazel-runfiles-root/setup.MODULE.bazel

@@ -0,0 +1,8 @@
+uv = use_extension("@aspect_rules_py//uv:extensions.bzl", "uv")
+uv.declare_hub(hub_name = "pypi_uv_bazel_runfiles_root")
+uv.project(
+    hub_name = "pypi_uv_bazel_runfiles_root",
+    lock = "//cases/uv-bazel-runfiles-root:uv.lock",
+    pyproject = "//cases/uv-bazel-runfiles-root:pyproject.toml",
+)
+use_repo(uv, "pypi_uv_bazel_runfiles_root")

e2e/cases/uv-bazel-runfiles-root/test_runfiles_root.py

@@ -0,0 +1,13 @@
+import os
+from pathlib import Path
+
+from runfiles import runfiles
+
+runfiles_dir = os.environ["RUNFILES_DIR"]
+assert runfiles._FindPythonRunfilesRoot() == runfiles_dir
+
+r = runfiles.Create()
+assert r is not None
+path = r.Rlocation("_main/cases/uv-bazel-runfiles-root/runfiles_root_data.txt")
+assert path is not None
+assert Path(path).read_text() == "runfiles root data\n"

py/private/py_venv/_virtualenv.py

@@ -11,6 +11,84 @@
     os.environ["PATH"] = _venv_bin + os.pathsep + os.environ.get("PATH", "")
 del _venv_bin
 
+
+def _bazel_runfiles_root_from_env():
+    runfiles_dir = os.environ.get("RUNFILES_DIR")
+    if runfiles_dir:
+        return runfiles_dir
+
+    manifest = os.environ.get("RUNFILES_MANIFEST_FILE")
+    if not manifest:
+        return None
+    if manifest.endswith("/MANIFEST") or manifest.endswith("\\MANIFEST"):
+        return manifest[: -len("/MANIFEST")]
+    if manifest.endswith(".runfiles_manifest"):
+        return manifest[: -len("_manifest")]
+    return None
+
+
+def _patch_bazel_runfiles_module(bazel_runfiles):
+    original = getattr(bazel_runfiles, "_FindPythonRunfilesRoot", None)
+    if original is None or getattr(original, "_rules_py_env_patch", False):
+        return
+
+    def _find_python_runfiles_root():
+        return _bazel_runfiles_root_from_env() or original()
+
+    _find_python_runfiles_root._rules_py_env_patch = True
+    bazel_runfiles._FindPythonRunfilesRoot = _find_python_runfiles_root
+
+
+class _BazelRunfilesFinder:
+    fullname = None
+
+    def find_spec(self, fullname, path, target=None):  # noqa: ARG002
+        if fullname != "runfiles.runfiles" or self.fullname is not None:
+            return None
+
+        from functools import partial
+        from importlib.util import find_spec
+
+        self.fullname = fullname
+        try:
+            spec = find_spec(fullname, path)
+            if spec is None or spec.loader is None:
+                return spec
+            func_name = "exec_module" if hasattr(spec.loader, "exec_module") else "load_module"
+            old = getattr(spec.loader, func_name)
+            func = self.exec_module if func_name == "exec_module" else self.load_module
+            if old is not func:
+                setattr(spec.loader, func_name, partial(func, old))
+            return spec
+        finally:
+            self.fullname = None
+
+    @staticmethod
+    def exec_module(old, module):
+        old(module)
+        _patch_bazel_runfiles_module(module)
+
+    @staticmethod
+    def load_module(old, name):
+        module = old(name)
+        _patch_bazel_runfiles_module(module)
+        return module
+
+
+def _patch_bazel_runfiles_root():
+    if not _bazel_runfiles_root_from_env():
+        return
+
+    module = sys.modules.get("runfiles.runfiles")
+    if module is not None:
+        _patch_bazel_runfiles_module(module)
+        return
+
+    sys.meta_path.insert(0, _BazelRunfilesFinder())
+
+
+_patch_bazel_runfiles_root()
+
 VIRTUALENV_PATCH_FILE = os.path.join(__file__)