fix (venv_site_packages): Fix wrong runfiles.symlinks when py_binary is not in root module (#3505)

When: 1) venv_site_packages is on 2) we have a py_binary in a non-root
module (e.g. a tool used in rules), and it uses python deps that result
in usage of runfiles.symlinks, the `ctx.runfile` based symlinks end up
going in the wrong folder (`_main`), while the
`ctx.actions.symlink(...)` files go in the right place. This results in
an invalid `venv` and import errors.

The reason is that `actions.symlinks` always go to the `_main` module,
as [Bazel docs
explain](https://bazel.build/extending/rules#runfiles_symlinks). To send
symlinks to other modules, one needs to use root_symlinks and prefix
them with the right module name.

Fixes: https://github.com/bazel-contrib/rules_python/issues/3503

---------

Co-authored-by: Shayan Hoshyari <hoshyari@adobe.com>
Co-authored-by: Richard Levasseur <rlevasseur@google.com>

python/private/py_executable.bzl

diff --git a/python/private/py_executable.bzl b/python/private/py_executable.bzl
index 1b884e9..5bac624 100644
--- a/python/private/py_executable.bzl
+++ b/python/private/py_executable.bzl
@@ -629,8 +629,10 @@
         ),
         # venv files for user library dependencies (files that are specific
         # to the executable bootstrap and python runtime aren't here).
+        # `root_symlinks` should be used, otherwise, with symlinks files always go
+        # to `_main` prefix, and binaries from non-root module become broken.
         lib_runfiles = ctx.runfiles(
-            symlinks = venv_app_files.runfiles_symlinks,
+            root_symlinks = venv_app_files.runfiles_symlinks,
         ),
     )
 

python/private/venv_runfiles.bzl

diff --git a/python/private/venv_runfiles.bzl b/python/private/venv_runfiles.bzl
index 02b18e5..7f6af0c 100644
--- a/python/private/venv_runfiles.bzl
+++ b/python/private/venv_runfiles.bzl
@@ -65,7 +65,10 @@
             bin_venv_path = paths.join(base, venv_path)
             if is_file(link_to):
                 # use paths.join to handle ctx.label.package = ""
-                symlink_from = paths.join(ctx.label.package, bin_venv_path)
+                # runfile_prefix should be prepended as we use runfiles.root_symlinks
+                runfile_prefix = ctx.label.repo_name or ctx.workspace_name
+                symlink_from = paths.join(runfile_prefix, ctx.label.package, bin_venv_path)
+
                 runfiles_symlinks[symlink_from] = link_to
             else:
                 venv_link = ctx.actions.declare_symlink(bin_venv_path)

tests/modules/other/BUILD.bazel

diff --git a/tests/modules/other/BUILD.bazel b/tests/modules/other/BUILD.bazel
index 46f1b96..665049b 100644
--- a/tests/modules/other/BUILD.bazel
+++ b/tests/modules/other/BUILD.bazel
@@ -1,3 +1,4 @@
+load("@rules_python//python:py_binary.bzl", "py_binary")
 load("@rules_python//tests/support:py_reconfig.bzl", "py_reconfig_binary")
 
 package(
@@ -12,3 +13,18 @@
     bootstrap_impl = "system_python",
     main = "external_main.py",
 )
+
+py_binary(
+    name = "venv_bin",
+    srcs = ["venv_bin.py"],
+    config_settings = {
+        "@rules_python//python/config_settings:bootstrap_impl": "script",
+        "@rules_python//python/config_settings:venvs_site_packages": "yes",
+    },
+    deps = [
+        # Add two packages that install into the same directory. This is
+        # to test that namespace packages install correctly and are importable.
+        "//nspkg_delta",
+        "//nspkg_gamma",
+    ],
+)

tests/modules/other/venv_bin.py

diff --git a/tests/modules/other/venv_bin.py b/tests/modules/other/venv_bin.py
new file mode 100644
index 0000000..5455b23
--- /dev/null
+++ b/tests/modules/other/venv_bin.py
@@ -0,0 +1,17 @@
+import nspkg
+
+print(nspkg)
+
+import nspkg.subnspkg
+
+print(nspkg.subnspkg)
+
+import nspkg.subnspkg.delta
+
+print(nspkg.subnspkg.delta)
+
+import nspkg.subnspkg.gamma
+
+print(nspkg.subnspkg.gamma)
+
+print("@other//:venv_bin ran successfully.")

tests/venv_site_packages_libs/BUILD.bazel

diff --git a/tests/venv_site_packages_libs/BUILD.bazel b/tests/venv_site_packages_libs/BUILD.bazel
index 2eb9678..e573dc6 100644
--- a/tests/venv_site_packages_libs/BUILD.bazel
+++ b/tests/venv_site_packages_libs/BUILD.bazel
@@ -1,3 +1,4 @@
+load("@rules_shell//shell:sh_test.bzl", "sh_test")
 load("//python:py_library.bzl", "py_library")
 load("//tests/support:py_reconfig.bzl", "py_reconfig_test")
 load(
@@ -19,6 +20,20 @@
     ],
 )
 
+sh_test(
+    name = "py_binary_other_module_test",
+    srcs = [
+        "py_binary_other_module_test.sh",
+    ],
+    data = [
+        "@other//:venv_bin",
+    ],
+    env = {
+        "VENV_BIN": "$(rootpath @other//:venv_bin)",
+    },
+    target_compatible_with = NOT_WINDOWS,
+)
+
 py_reconfig_test(
     name = "venvs_site_packages_libs_test",
     srcs = ["bin.py"],

tests/venv_site_packages_libs/py_binary_other_module_test.sh

diff --git a/tests/venv_site_packages_libs/py_binary_other_module_test.sh b/tests/venv_site_packages_libs/py_binary_other_module_test.sh
new file mode 100755
index 0000000..9fb71d5
--- /dev/null
+++ b/tests/venv_site_packages_libs/py_binary_other_module_test.sh
@@ -0,0 +1,6 @@
+# Test that for a py_binary from a dependency module, we place links created via
+# runfiles(...) in the right place. This tests the fix made for issues/3503
+
+set -eu
+echo "[*] Testing running the binary"
+"$VENV_BIN"