wip: add symlinks for .dist-info folders in the venv

This adds the necessary `.dist-info` files into the mix
and should get us reasonably close to handling 99% of the cases.

The expected differences from this and a `venv` built by `uv` would be:
* Shared libraries are stored in `<package>.libs` in `uv` venvs. This
  can be achieved in `rules_python` by changing the `installer` settings
  in the `wheel_installer/wheel.py#unzip` function.
* The `RECORD` files are excluded from the `venv`s for better cache hit
  rate in `bazel`, however I am not sure if we should do that for actual
  wheels that are downloaded from the internet.

Tested:
- [x] Building the `//docs` and manually checking the symlinks.
- [ ] Unit tests

Work towards #2156

Author: aignas

python/private/BUILD.bazel

@@ -444,6 +444,7 @@ bzl_library(
         ":attributes_bzl",
         ":common_bzl",
         ":flags_bzl",
+        ":normalize_name_bzl",
         ":precompile_bzl",
         ":py_cc_link_params_info_bzl",
         ":py_internal_bzl",

python/private/common.bzl

@@ -396,8 +396,8 @@ 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
-            `(runfiles_path, site_packages_path)` for symlinks to create
+        site_packages_symlinks: {type}`list[tuple[str, str, str]]` tuples of
+            `(normalized_package_name, runfiles_path, site_packages_path)` for symlinks to create
             in the consuming binary's venv site packages.
 
     Returns:

python/private/py_executable.bzl

@@ -671,11 +671,19 @@ def _create_site_packages_symlinks(ctx, site_packages):
 
 def _build_link_map(entries):
     link_map = {}
-    for link_to_runfiles_path, site_packages_path in entries:
-        if site_packages_path in link_map:
+    package_map = {}
+    for link_to_runfiles_path, site_packages_path, package in entries:
+        if package and package in package_map:
             # We ignore duplicates by design. The dependency closer to the
             # binary gets precedence due to the topological ordering.
             continue
+        elif site_packages_path in link_map:
+            # The packages are not duplicate, but have overlapping paths, for
+            # now we ignore these.
+            continue
+        elif package:
+            package_map[package] = site_packages_path
+            link_map[site_packages_path] = link_to_runfiles_path
         else:
             link_map[site_packages_path] = link_to_runfiles_path
 

python/private/py_info.bzl

@@ -148,11 +148,11 @@ The order of the depset is not guaranteed and may be changed in the future. It
 is recommended to use `default` order (the default).
 """,
         "site_packages_symlinks": """
-:type: depset[tuple[str | None, str]]
+:type: depset[tuple[str | None, str, str | None]]
 
 A depset with `topological` ordering.
 
-Tuples of `(runfiles_path, site_packages_path)`. Where
+Tuples of `(runfiles_path, site_packages_path, symlink_key)`. Where
 * `runfiles_path` is a runfiles-root relative path. It is the path that
   has the code to make importable. If `None` or empty string, then it means
   to not create a site packages directory with the `site_packages_path`
@@ -161,6 +161,11 @@ Tuples of `(runfiles_path, site_packages_path)`. Where
   the venv for whatever creates the venv (typically py_binary). It makes
   the code in `runfiles_path` available for import. Note that this
   is created as a "raw" symlink (via `declare_symlink`).
+* `symlink_key` is a path which should be checked instead of the
+  `site_packages_path`. This is typically used for dist-info directories as
+  they include version in the filenames and we want to ensure that we can
+  replace the whole package and not have orphan .dist-info packages in the
+  virtual environment.
 
 :::{include} /_includes/experimental_api.md
 :::

python/private/py_library.bzl

@@ -41,6 +41,7 @@ load(
     "runfiles_root_path",
 )
 load(":flags.bzl", "AddSrcsToRunfilesFlag", "PrecompileFlag", "VenvsSitePackages")
+load(":normalize_name.bzl", "normalize_name")
 load(":precompile.bzl", "maybe_precompile")
 load(":py_cc_link_params_info.bzl", "PyCcLinkParamsInfo")
 load(":py_info.bzl", "PyInfo")
@@ -206,16 +207,32 @@ Source files are no longer added to the runfiles directly.
 :::
 """
 
+def _get_distinfo_metadata(ctx):
+    data = ctx.files.data or []
+    for d in data:
+        # work on case insensitive FSes
+        if d.basename.lower() != "metadata":
+            continue
+
+        if d.dirname.endswith(".dist-info"):
+            return d
+
+    return None
+
 def _get_imports_and_site_packages_symlinks(ctx, semantics):
     imports = depset()
     site_packages_symlinks = depset()
     if VenvsSitePackages.is_enabled(ctx):
-        site_packages_symlinks = _get_site_packages_symlinks(ctx)
+        dist_info_metadata = _get_distinfo_metadata(ctx)
+        site_packages_symlinks = _get_site_packages_symlinks(
+            ctx,
+            dist_info_metadata,
+        )
     else:
         imports = collect_imports(ctx, semantics)
     return imports, site_packages_symlinks
 
-def _get_site_packages_symlinks(ctx):
+def _get_site_packages_symlinks(ctx, dist_info_metadata):
     imports = ctx.attr.imports
     if len(imports) == 0:
         fail("When venvs_site_packages is enabled, exactly one `imports` " +
@@ -254,6 +271,23 @@ def _get_site_packages_symlinks(ctx):
     repo_runfiles_dirname = None
     dirs_with_init = {}  # dirname -> runfile path
     site_packages_symlinks = []
+    if dist_info_metadata:
+        # in order to be able to have replacements in the venv, we have to add a
+        # third value into the site_packages_symlinks, which would be the normalized
+        # package name. This allows us to ensure that we can replace the `dist-info`
+        # directories by checking if the package key is there.
+        dist_info_dir = paths.basename(dist_info_metadata.dirname)
+        package, _, _suffix = dist_info_dir.rpartition(".dist-info")
+        package, _, _version = package.rpartition("-")
+        symlink_key = "{}.dist-info".format(normalize_name(package))
+
+        repo_runfiles_dirname = runfiles_root_path(ctx, dist_info_metadata.short_path).partition("/")[0]
+        site_packages_symlinks.append((
+            paths.join(repo_runfiles_dirname, site_packages_root, dist_info_dir),
+            dist_info_dir,
+            symlink_key,
+        ))
+
     for src in ctx.files.srcs:
         if src.extension not in PYTHON_FILE_EXTENSIONS:
             continue
@@ -274,6 +308,7 @@ def _get_site_packages_symlinks(ctx):
             site_packages_symlinks.append((
                 paths.join(repo_runfiles_dirname, site_packages_root, filename),
                 filename,
+                None,
             ))
 
     # Sort so that we encounter `foo` before `foo/bar`. This ensures we
@@ -294,6 +329,7 @@ def _get_site_packages_symlinks(ctx):
         site_packages_symlinks.append((
             paths.join(repo_runfiles_dirname, site_packages_root, dirname),
             dirname,
+            None,
         ))
     return site_packages_symlinks