fix: write __init__ shim for pip runfiles helper (#519)

Fixes #377 

Even though messing with __init__.py is a little risky, in this case its
safer since we only generate it while symlinking files from external
repositories into venv, and since
https://github.com/bazelbuild/rules_python/blob/18f76f9cc4c3c6c4470f7a881e6ea4e8b80d7bab/python/runfiles/__init__.py#L15
is simple line, i can't think of easier fix than this.

### Changes are visible to end-users: no

- Searched for relevant documentation and updated as needed: yes
- Breaking change (forces users to change their own code or config): no
- Suggested release notes appear below: no

### Test plan

- New test cases added

Author: thesayyn

gazelle_python.yaml

@@ -3833,6 +3833,8 @@ manifest:
     rfc3986_validator: rfc3986_validator
     rpds: rpds_py
     rpds.rpds: rpds_py
+    runfiles: bazel_runfiles
+    runfiles.runfiles: bazel_runfiles
     s3transfer: s3transfer
     s3transfer.bandwidth: s3transfer
     s3transfer.compat: s3transfer
@@ -4030,4 +4032,4 @@ manifest:
     yaml.tokens: PyYAML
   pip_repository:
     name: pypi
-integrity: 6359bc010056b6e19ee6cbb2fc9a8ac841449521e3b4997fec58163e5c635496
+integrity: 7049b81983381a0118bd63ae8cb7f19126d0034fcbb256281453ee5cbf3499be

py/tests/py-binary/BUILD.bazel

@@ -1,6 +1,7 @@
 load("@aspect_bazel_lib//lib:run_binary.bzl", "run_binary")
 load("@aspect_bazel_lib//lib:testing.bzl", "assert_contains")
-load("//py:defs.bzl", "py_binary")
+load("@pypi//:requirements.bzl", "requirement")
+load("//py:defs.bzl", "py_binary", "py_test")
 
 #################
 # Case 1: main can be a string referencing some source file
@@ -48,3 +49,13 @@ assert_contains(
     actual = "main_from_name.out",
     expected = "running main_from_name",
 )
+
+#################
+# Case 3: Runfiles from pip works
+py_test(
+    name = "runfiles_from_pip_test",
+    srcs = ["runfiles_from_pip.py"],
+    data = ["test_data.txt"],
+    main = "runfiles_from_pip.py",
+    deps = [requirement("bazel-runfiles")],
+)

py/tests/py-binary/runfiles_from_pip.py

@@ -0,0 +1,14 @@
+import os
+import pathlib
+import unittest
+
+import runfiles  # requirement("bazel-runfiles")
+
+class RunfilesTest(unittest.TestCase):
+    def test_runfiles(self) -> None:
+        r = runfiles.Runfiles.Create()
+        path = pathlib.Path(r.Rlocation(os.getenv("BAZEL_WORKSPACE")+"/py/tests/py-binary/test_data.txt"))
+        self.assertEquals(path.read_text(), "42\n")
+
+if __name__ == "__main__":
+    unittest.main()

py/tests/py-binary/test_data.txt

@@ -0,0 +1 @@
+42

py/tools/py/src/pth.rs

@@ -6,6 +6,29 @@ use std::{
 
 use miette::{miette, Context, IntoDiagnostic, LabeledSpan, MietteDiagnostic, Severity};
 
+const RULES_PYTHON_INIT_PATH: &str = "runfiles/__init__.py";
+const RULES_PYTHON_RUNFILES_INIT_SHIM: &str = r#"
+# Generated by rules_py venv for rules_python compatibility
+# See: https://github.com/bazelbuild/rules_python/pull/2115
+# See: https://github.com/aspect-build/rules_py/issues/377
+from . import runfiles
+def _FindPythonRunfilesRoot():
+    root = __file__
+    # The original implementation of this function in rules_python expects the runfiles root to be 4 directories up from the current file.
+    # but in rules_py there is additional two segments that it needs to go up to  find the runfiles root.
+    # bazel-bin/py/tests/external-deps/foo.runfiles/.foo.venv/lib/python3.9/site-packages/runfiles
+    # ╰─────────────────────┬─────────────────────╯╰───┬───╯╰─────────────┬─────────────╯╰───┬───╯
+    #             bazel runfiles root              venv root    Python packages root   Python package
+    
+    for _ in range("rules_python/python/runfiles/runfiles.py".count("/") + 3):
+        root = os.path.dirname(root)
+    return root
+
+runfiles._FindPythonRunfilesRoot = _FindPythonRunfilesRoot
+
+from .runfiles import *
+"#;
+
 /// Strategy that will be used when creating the virtual env symlink and
 /// a collision is found.
 #[derive(Default, Debug)]
@@ -121,6 +144,19 @@ fn create_symlinks(
         // itself is not a regular package and is not supposed to have an `__init__.py` file.
         if path.is_dir() {
             create_symlinks(&path, root_dir, dst_dir, collision_strategy)?;
+        }
+        // rules_python runfiles helper needs some special handling when consumed as pip dependency.
+        // See: https://github.com/aspect-build/rules_py/issues/377
+        // See: https://github.com/bazelbuild/rules_python/pull/2115
+        else if path.as_path().ends_with(RULES_PYTHON_INIT_PATH) {
+            // Instead of symlinking the __init__.py file from its original location to the venv site-packages,
+            // we write a shim __init__.py that makes the runfiles pypi library work with rules_py.
+            fs::write(
+                dst_dir.join(RULES_PYTHON_INIT_PATH),
+                RULES_PYTHON_RUNFILES_INIT_SHIM,
+            )
+            .into_diagnostic()
+            .wrap_err("Unable to write to runfiles __init__.py file")?;
         } else if dir != root_dir || entry.file_name() != "__init__.py" {
             create_symlink(&entry, root_dir, dst_dir, collision_strategy)?;
         }

requirements.in

@@ -8,3 +8,4 @@ snakesay
 ftfy==6.2.0
 neptune==1.10.2
 six
+bazel-runfiles

requirements.txt

@@ -18,6 +18,9 @@ attrs==23.2.0 \
     # via
     #   jsonschema
     #   referencing
+bazel-runfiles==1.1.0 \
+    --hash=sha256:37f59ea505b86ada391ef94e0949ff38a6fd6c111c9a8338065b16b355d0efae
+    # via -r requirements.in
 boto3==1.34.93 \
     --hash=sha256:b59355bf4a1408563969526f314611dbeacc151cf90ecb22af295dcc4fe18def \
     --hash=sha256:e39516e4ca21612932599819662759c04485d53ca457996a913163da11f052a4