Dynamic resolution of runfiles-relative import roots

Author: arrdem

py/private/py_venv/py_venv.bzl

@@ -57,19 +57,12 @@ def _py_venv_base_impl(ctx):
 
     # Check for duplicate virtual dependency names. Those that map to the same resolution target would have been merged by the depset for us.
     virtual_resolution = _py_library.resolve_virtuals(ctx)
+    # Note that this adds the workspace root for us (sigh), don't need to add to it
     imports_depset = _py_library.make_imports_depset(ctx, extra_imports_depsets = virtual_resolution.imports)
 
     pth_lines = ctx.actions.args()
     pth_lines.use_param_file("%s", use_always = True)
     pth_lines.set_param_file_format("multiline")
-
-    # FIXME: This was hardcoded in the original rule_py venv and is preserved
-    # for compatibility. Repo-absolute imports are Bad (TM) and shouldn't be on
-    # by default. I believe that as of recent rules_python, creating these
-    # repo-absolute imports is handled as part of the PyInfo calculation. If we
-    # get this from rules_python, it should be removed. Or it should be moved so
-    # that we calculate it as part of the imports depset logic.
-    pth_lines.add(".")
     pth_lines.add_all(imports_depset)
 
     site_packages_pth_file = ctx.actions.declare_file("{}.pth".format(ctx.attr.name))

py/tests/py_venv_conflict/test_import_roots.py

@@ -1,5 +1,12 @@
 #!/usr/bin/env python3
 
+import os
+for k, v in os.environ.items():
+    if k.startswith("BUILD_") or k.startswith("RUNFILES_"):
+        print(k, ":", v)
+
+print("---")
+
 from pathlib import Path
 
 # prefix components:

py/tools/py/BUILD.bazel

@@ -10,6 +10,7 @@ rust_library(
     ],
     data = [
         "src/_virtualenv.py",
+        "src/_aspect.py",
         "src/activate.tmpl",
         "src/pyvenv.cfg.tmpl",
         "src/runfiles_interpreter.tmpl",

py/tools/py/src/_aspect.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+
+_repo_mapping = {}
+
+def pathadd(path: str, workspace=""):
+    """Attempt to soundly append a path to the sys.path.
+
+    See https://github.com/aspect-build/rules_py/issues/573 for gorey details.
+
+    Because of Bazel's symlinking behavior and `exec()` file canonicalization,
+    the `sys.prefix` and site dir may escape from `.runfiles/` as created by
+    Bazel. This is a problem for a ton of reasons, notably that source files
+    aren't copied to `bazel-bin`, so relative paths from the (canonicalized!)
+    site dir back "up" to firstparty sources may be broken.
+
+    Yay.
+
+    As a workaround, this codepath exists to try and use `$RUNFILES_DIR` (either
+    set by Bazel or enforced by the Aspect venv `activate` script) as the basis
+    for appending runfiles-relative logical paths to the `sys.path`.
+
+    """
+
+    runfiles_dir = os.getenv("RUNFILES_DIR")
+    if not runfiles_dir:
+        print("{} ERROR: Unable to identify the runfiles root!".format(__file__))
+        exit(1)
+
+    # Parse the repo mapping file if it exists
+    if not _repo_mapping:
+        repo_mapping_file = os.path.join(runfiles_dir, "_repo_mapping")
+        if os.path.exists(repo_mapping_file):
+            with open(repo_mapping_file, "r") as fp:
+                for line in fp:
+                    line = line.strip()
+                    c, a, mca = line.split(",")
+                    _repo_mapping[(c, a)] = mca
+
+    # Deal with looking up a remapped repo name
+    # HACK: This assumes the built binary is in the main workspace
+    if _repo_mapping:
+        repo, path = path.split("/", 1)
+        repo = _repo_mapping[(workspace, repo)]
+        path = os.path.join(repo, path)
+        
+    path = os.path.normpath(os.path.join(runfiles_dir, path))
+    if path not in sys.path:
+        sys.path.append(path)

py/tools/py/src/runfiles_interpreter.tmpl

@@ -80,7 +80,7 @@ source "${RUNFILES_DIR:-/dev/null}/${f}" 2>/dev/null || \
 } >/dev/null
 
 # Look up the runfiles-based interpreter and put its dir _first_ on the path.
-INTERPRETER="$(realpath $(rlocation {{INTERPRETER_TARGET}}))"
+INTERPRETER="$(rlocation {{INTERPRETER_TARGET}})"
 
 # Figure out if we're dealing with just some program or a real install
 # <SOMEDIR> <- possible $PYTHONHOME

py/tools/py/src/venv.rs

@@ -380,6 +380,13 @@ pub fn create_empty_venv<'a>(
     )
     .into_diagnostic()?;
 
+    // And include our horrid path machinery
+    fs::write(
+        &venv.site_dir.join("_aspect.py"),
+        include_str!("_aspect.py"),
+    )
+    .into_diagnostic()?;
+
     Ok(venv)
 }
 
@@ -441,10 +448,11 @@ pub fn populate_venv_with_copies(
     let action_bin_dir = action_src_dir.join(bin_dir);
 
     #[cfg(feature = "debug")]
-    eprintln!("action_src_dir: {}", &action_src_dir.to_str().unwrap());
-
-    #[cfg(feature = "debug")]
-    eprintln!("action_bin_dir: {}", &action_bin_dir.to_str().unwrap());
+    {
+        println!("action_src_dir: {}", &action_src_dir.to_str().unwrap());
+        println!("action_bin_dir: {}", &action_bin_dir.to_str().unwrap());
+        println!("venv.site_dir: {}", &venv.site_dir.to_str().unwrap());
+    }
 
     let source_pth = File::open(pth_file.src.as_path())
         .into_diagnostic()
@@ -459,7 +467,7 @@ pub fn populate_venv_with_copies(
         .write(
             b"\
 # Generated by Aspect py_binary
-# Contains relative import paths to non site-package trees within the .runfiles
+# Sets up imports within the .runfiles tree
 ",
         )
         .into_diagnostic()?;
@@ -532,14 +540,16 @@ pub fn populate_venv_with_copies(
             //
             // [1] https://github.com/python/cpython/blob/ce31ae5209c976d28d1c21fcbb06c0ae5e50a896/Lib/site.py#L215
 
-            let path_to_runfiles =
-                diff_paths(&action_bin_dir, action_bin_dir.join(&venv.site_dir)).unwrap();
-
             writeln!(dest_pth_writer, "# @{}", line).into_diagnostic()?;
             writeln!(
                 dest_pth_writer,
-                "{}",
-                path_to_runfiles.join(entry).to_str().unwrap()
+                "import _aspect;_aspect.pathadd('{}', workspace='{}')",
+                line.replace("'", "\\'"),
+                if workspace == main_module {
+                    &""
+                } else {
+                    main_module.to_str().unwrap()
+                },
             )
             .into_diagnostic()?;
         }

py/tools/venv_shim/src/main.rs

@@ -64,7 +64,7 @@ fn find_python_executables(version_from_cfg: &str, exclude_dir: &Path) -> Option
                 None
             }
         })
-        .filter_map(|path| path.canonicalize().ok())
+        //.filter_map(|path| path.canonicalize().ok())
         .filter(|potential_executable| potential_executable.parent() != Some(exclude_dir))
         .filter(|potential_executable| compare_versions(version_from_cfg, &potential_executable))
         .collect();