refactor: change site_packages_symlinks to venv_symlinks (#2939)

This generalizes the ability to populate the venv directory by adding
and additional field,
`kind`, which tells which directory of the venv to populate. A symbolic
constant is used
to indicate which directory so that users don't have to re-derive the
platform and version
specific paths that make up the venv directory names.

This follows the design described by
#2156 (comment)

This also changes it to a depset of structs to make it more forward
compatible. A provider
is used because they're slightly more memory efficient than regular
structs.

Work towards #2156

Author: rickeylev

CHANGELOG.md

@@ -70,6 +70,8 @@ END_UNRELEASED_TEMPLATE
   `_test` target is deprecated and will be removed in the next major release.
   ([#2794](https://github.com/bazel-contrib/rules_python/issues/2794)
 * (py_wheel) py_wheel always creates zip64-capable wheel zips
+* (providers) (experimental) {obj}`PyInfo.venv_symlinks` replaces
+  `PyInfo.site_packages_symlinks`
 
 {#v0-0-0-fixed}
 ### Fixed
@@ -203,7 +205,7 @@ END_UNRELEASED_TEMPLATE
   please check the {obj}`uv.configure` tag class.
 * Add support for riscv64 linux platform.
 * (toolchains) Add python 3.13.2 and 3.12.9 toolchains
-* (providers) (experimental) {obj}`PyInfo.site_packages_symlinks` field added to
+* (providers) (experimental) `PyInfo.site_packages_symlinks` field added to
   allow specifying links to create within the venv site packages (only
   applicable with {obj}`--bootstrap_impl=script`)
   ([#2156](https://github.com/bazelbuild/rules_python/issues/2156)).

python/features.bzl

@@ -31,11 +31,11 @@ def _features_typedef():
     :::
     ::::
 
-    ::::{field} py_info_site_packages_symlinks
+    ::::{field} py_info_venv_symlinks
 
-    True if the `PyInfo.site_packages_symlinks` field is available.
+    True if the `PyInfo.venv_symlinks` field is available.
 
-    :::{versionadded} 1.4.0
+    :::{versionadded} VERSION_NEXT_FEATURE
     :::
     ::::
 
@@ -61,7 +61,7 @@ features = struct(
     TYPEDEF = _features_typedef,
     # keep sorted
     precompile = True,
-    py_info_site_packages_symlinks = True,
+    py_info_venv_symlinks = True,
     uses_builtin_rules = not config.enable_pystar,
     version = _VERSION_PRIVATE if "$Format" not in _VERSION_PRIVATE else "",
 )

python/private/attributes.bzl

@@ -260,7 +260,7 @@ The order of this list can matter because it affects the order that information
 from dependencies is merged in, which can be relevant depending on the ordering
 mode of depsets that are merged.
 
-* {obj}`PyInfo.site_packages_symlinks` uses topological ordering.
+* {obj}`PyInfo.venv_symlinks` uses topological ordering.
 
 See {obj}`PyInfo` for more information about the ordering of its depsets and
 how its fields are merged.

python/private/common.bzl

@@ -378,7 +378,7 @@ def create_py_info(
         implicit_pyc_files,
         implicit_pyc_source_files,
         imports,
-        site_packages_symlinks = []):
+        venv_symlinks = []):
     """Create PyInfo provider.
 
     Args:
@@ -396,7 +396,7 @@ def create_py_info(
         implicit_pyc_files: {type}`depset[File]` Implicitly generated pyc files
             that a binary can choose to include.
         imports: depset of strings; the import path values to propagate.
-        site_packages_symlinks: {type}`list[tuple[str, str]]` tuples of
+        venv_symlinks: {type}`list[tuple[str, str]]` tuples of
             `(runfiles_path, site_packages_path)` for symlinks to create
             in the consuming binary's venv site packages.
 
@@ -406,7 +406,7 @@ def create_py_info(
         necessary for deprecated extra actions support).
     """
     py_info = PyInfoBuilder.new()
-    py_info.site_packages_symlinks.add(site_packages_symlinks)
+    py_info.venv_symlinks.add(venv_symlinks)
     py_info.direct_original_sources.add(original_sources)
     py_info.direct_pyc_files.add(required_pyc_files)
     py_info.direct_pyi_files.add(ctx.files.pyi_srcs)

python/private/flags.bzl

@@ -154,12 +154,12 @@ def _venvs_site_packages_is_enabled(ctx):
     flag_value = ctx.attr.experimental_venvs_site_packages[BuildSettingInfo].value
     return flag_value == VenvsSitePackages.YES
 
-# Decides if libraries try to use a site-packages layout using site_packages_symlinks
+# Decides if libraries try to use a site-packages layout using venv_symlinks
 # buildifier: disable=name-conventions
 VenvsSitePackages = FlagEnum(
-    # Use site_packages_symlinks
+    # Use venv_symlinks
     YES = "yes",
-    # Don't use site_packages_symlinks
+    # Don't use venv_symlinks
     NO = "no",
     is_enabled = _venvs_site_packages_is_enabled,
 )

python/private/py_executable.bzl

@@ -54,7 +54,7 @@ load(":flags.bzl", "BootstrapImplFlag", "VenvsUseDeclareSymlinkFlag")
 load(":precompile.bzl", "maybe_precompile")
 load(":py_cc_link_params_info.bzl", "PyCcLinkParamsInfo")
 load(":py_executable_info.bzl", "PyExecutableInfo")
-load(":py_info.bzl", "PyInfo")
+load(":py_info.bzl", "PyInfo", "VenvSymlinkKind")
 load(":py_internal.bzl", "py_internal")
 load(":py_runtime_info.bzl", "DEFAULT_STUB_SHEBANG", "PyRuntimeInfo")
 load(":reexports.bzl", "BuiltinPyInfo", "BuiltinPyRuntimeInfo")
@@ -543,6 +543,7 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         VenvsUseDeclareSymlinkFlag.get_value(ctx) == VenvsUseDeclareSymlinkFlag.YES
     )
     recreate_venv_at_runtime = False
+    bin_dir = "{}/bin".format(venv)
 
     if not venvs_use_declare_symlink_enabled or not runtime.supports_build_time_venv:
         recreate_venv_at_runtime = True
@@ -556,7 +557,7 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         # When the venv symlinks are disabled, the $venv/bin/python3 file isn't
         # needed or used at runtime. However, the zip code uses the interpreter
         # File object to figure out some paths.
-        interpreter = ctx.actions.declare_file("{}/bin/{}".format(venv, py_exe_basename))
+        interpreter = ctx.actions.declare_file("{}/{}".format(bin_dir, py_exe_basename))
         ctx.actions.write(interpreter, "actual:{}".format(interpreter_actual_path))
 
     elif runtime.interpreter:
@@ -568,7 +569,7 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         # declare_symlink() is required to ensure that the resulting file
         # in runfiles is always a symlink. An RBE implementation, for example,
         # may choose to write what symlink() points to instead.
-        interpreter = ctx.actions.declare_symlink("{}/bin/{}".format(venv, py_exe_basename))
+        interpreter = ctx.actions.declare_symlink("{}/{}".format(bin_dir, py_exe_basename))
 
         interpreter_actual_path = runfiles_root_path(ctx, runtime.interpreter.short_path)
         rel_path = relative_path(
@@ -581,7 +582,7 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         ctx.actions.symlink(output = interpreter, target_path = rel_path)
     else:
         py_exe_basename = paths.basename(runtime.interpreter_path)
-        interpreter = ctx.actions.declare_symlink("{}/bin/{}".format(venv, py_exe_basename))
+        interpreter = ctx.actions.declare_symlink("{}/{}".format(bin_dir, py_exe_basename))
         ctx.actions.symlink(output = interpreter, target_path = runtime.interpreter_path)
         interpreter_actual_path = runtime.interpreter_path
 
@@ -618,89 +619,104 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         },
         computed_substitutions = computed_subs,
     )
-    site_packages_symlinks = _create_site_packages_symlinks(ctx, site_packages)
+
+    venv_dir_map = {
+        VenvSymlinkKind.BIN: bin_dir,
+        VenvSymlinkKind.LIB: site_packages,
+    }
+    venv_symlinks = _create_venv_symlinks(ctx, venv_dir_map)
 
     return struct(
         interpreter = interpreter,
         recreate_venv_at_runtime = recreate_venv_at_runtime,
         # Runfiles root relative path or absolute path
         interpreter_actual_path = interpreter_actual_path,
-        files_without_interpreter = [pyvenv_cfg, pth, site_init] + site_packages_symlinks,
+        files_without_interpreter = [pyvenv_cfg, pth, site_init] + venv_symlinks,
         # string; venv-relative path to the site-packages directory.
         venv_site_packages = venv_site_packages,
     )
 
-def _create_site_packages_symlinks(ctx, site_packages):
-    """Creates symlinks within site-packages.
+def _create_venv_symlinks(ctx, venv_dir_map):
+    """Creates symlinks within the venv.
 
     Args:
         ctx: current rule ctx
-        site_packages: runfiles-root-relative path to the site-packages directory
+        venv_dir_map: mapping of VenvSymlinkKind constants to the
+            venv path.
 
     Returns:
         {type}`list[File]` list of the File symlink objects created.
     """
 
-    # maps site-package symlink to the runfiles path it should point to
+    # maps venv-relative path to the runfiles path it should point to
     entries = depset(
         # NOTE: Topological ordering is used so that dependencies closer to the
         # binary have precedence in creating their symlinks. This allows the
         # binary a modicum of control over the result.
         order = "topological",
         transitive = [
-            dep[PyInfo].site_packages_symlinks
+            dep[PyInfo].venv_symlinks
             for dep in ctx.attr.deps
             if PyInfo in dep
         ],
     ).to_list()
+
     link_map = _build_link_map(entries)
+    venv_files = []
+    for kind, kind_map in link_map.items():
+        base = venv_dir_map[kind]
+        for venv_path, link_to in kind_map.items():
+            venv_link = ctx.actions.declare_symlink(paths.join(base, venv_path))
+            venv_link_rf_path = runfiles_root_path(ctx, venv_link.short_path)
+            rel_path = relative_path(
+                # dirname is necessary because a relative symlink is relative to
+                # the directory the symlink resides within.
+                from_ = paths.dirname(venv_link_rf_path),
+                to = link_to,
+            )
+            ctx.actions.symlink(output = venv_link, target_path = rel_path)
+            venv_files.append(venv_link)
 
-    sp_files = []
-    for sp_dir_path, link_to in link_map.items():
-        sp_link = ctx.actions.declare_symlink(paths.join(site_packages, sp_dir_path))
-        sp_link_rf_path = runfiles_root_path(ctx, sp_link.short_path)
-        rel_path = relative_path(
-            # dirname is necessary because a relative symlink is relative to
-            # the directory the symlink resides within.
-            from_ = paths.dirname(sp_link_rf_path),
-            to = link_to,
-        )
-        ctx.actions.symlink(output = sp_link, target_path = rel_path)
-        sp_files.append(sp_link)
-    return sp_files
+    return venv_files
 
 def _build_link_map(entries):
+    # dict[str kind, dict[str rel_path, str link_to_path]]
     link_map = {}
-    for link_to_runfiles_path, site_packages_path in entries:
-        if site_packages_path in link_map:
+    for entry in entries:
+        kind = entry.kind
+        kind_map = link_map.setdefault(kind, {})
+        if entry.venv_path in kind_map:
             # We ignore duplicates by design. The dependency closer to the
             # binary gets precedence due to the topological ordering.
             continue
         else:
-            link_map[site_packages_path] = link_to_runfiles_path
+            kind_map[entry.venv_path] = entry.link_to_path
 
     # An empty link_to value means to not create the site package symlink.
     # Because of the topological ordering, this allows binaries to remove
     # entries by having an earlier dependency produce empty link_to values.
-    for sp_dir_path, link_to in link_map.items():
-        if not link_to:
-            link_map.pop(sp_dir_path)
+    for kind, kind_map in link_map.items():
+        for dir_path, link_to in kind_map.items():
+            if not link_to:
+                kind_map.pop(dir_path)
 
-    # Remove entries that would be a child path of a created symlink.
-    # Earlier entries have precedence to match how exact matches are handled.
+    # dict[str kind, dict[str rel_path, str link_to_path]]
     keep_link_map = {}
-    for _ in range(len(link_map)):
-        if not link_map:
-            break
-        dirname, value = link_map.popitem()
-        keep_link_map[dirname] = value
-
-        prefix = dirname + "/"  # Add slash to prevent /X matching /XY
-        for maybe_suffix in link_map.keys():
-            maybe_suffix += "/"  # Add slash to prevent /X matching /XY
-            if maybe_suffix.startswith(prefix) or prefix.startswith(maybe_suffix):
-                link_map.pop(maybe_suffix)
 
+    # Remove entries that would be a child path of a created symlink.
+    # Earlier entries have precedence to match how exact matches are handled.
+    for kind, kind_map in link_map.items():
+        keep_kind_map = keep_link_map.setdefault(kind, {})
+        for _ in range(len(kind_map)):
+            if not kind_map:
+                break
+            dirname, value = kind_map.popitem()
+            keep_kind_map[dirname] = value
+            prefix = dirname + "/"  # Add slash to prevent /X matching /XY
+            for maybe_suffix in kind_map.keys():
+                maybe_suffix += "/"  # Add slash to prevent /X matching /XY
+                if maybe_suffix.startswith(prefix) or prefix.startswith(maybe_suffix):
+                    kind_map.pop(maybe_suffix)
     return keep_link_map
 
 def _map_each_identity(v):