feat: data and pyi files in the venv

.bazelrc

@@ -4,8 +4,8 @@
 # (Note, we cannot use `common --deleted_packages` because the bazel version command doesn't support it)
 # To update these lines, execute
 # `bazel run @rules_bazel_integration_test//tools:update_deleted_packages`
-build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/pythonconfig,gazelle/python/private,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered,tests/modules/other,tests/modules/other/nspkg_delta,tests/modules/other/nspkg_gamma,tests/modules/other/nspkg_single
-query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/pythonconfig,gazelle/python/private,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered,tests/modules/other,tests/modules/other/nspkg_delta,tests/modules/other/nspkg_gamma,tests/modules/other/nspkg_single
+build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/pythonconfig,gazelle/python/private,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered,tests/modules/other,tests/modules/other/nspkg_delta,tests/modules/other/nspkg_gamma,tests/modules/other/nspkg_single,tests/modules/other/simple_v1,tests/modules/other/simple_v2
+query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/patches,examples/bzlmod/py_proto_library,examples/bzlmod/py_proto_library/example.com/another_proto,examples/bzlmod/py_proto_library/example.com/proto,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/another_proto,examples/py_proto_library/example.com/proto,gazelle,gazelle/manifest,gazelle/manifest/generate,gazelle/manifest/hasher,gazelle/manifest/test,gazelle/modules_mapping,gazelle/python,gazelle/pythonconfig,gazelle/python/private,tests/integration/compile_pip_requirements,tests/integration/compile_pip_requirements_test_from_external_repo,tests/integration/custom_commands,tests/integration/ignore_root_user_error,tests/integration/ignore_root_user_error/submodule,tests/integration/local_toolchains,tests/integration/pip_parse,tests/integration/pip_parse/empty,tests/integration/py_cc_toolchain_registered,tests/modules/other,tests/modules/other/nspkg_delta,tests/modules/other/nspkg_gamma,tests/modules/other/nspkg_single,tests/modules/other/simple_v1,tests/modules/other/simple_v2
 
 test --test_output=errors
 

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,9 +396,9 @@ 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.
-        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.
+        venv_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.
 
     Returns:
         A tuple of the PyInfo instance and a depset of the

python/private/py_executable.bzl

@@ -680,17 +680,17 @@ def _create_venv_symlinks(ctx, venv_dir_map):
     return venv_files
 
 def _build_link_map(entries):
-    # dict[str kind, dict[str rel_path, str link_to_path]]
+    # dict[str kind, dict[str key or rel_path, tuple[str rel_path, str link_to_path]]]
     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:
-            kind_map[entry.venv_path] = entry.link_to_path
+
+        # We overwrite duplicates by design. The dependency closer to the
+        # binary gets precedence due to the topological ordering.
+        #
+        # This allows us to store only one version of the dist-info that is needed
+        kind_map[entry.key or entry.venv_path] = (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
@@ -710,7 +710,7 @@ def _build_link_map(entries):
         for _ in range(len(kind_map)):
             if not kind_map:
                 break
-            dirname, value = kind_map.popitem()
+            _, (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():

python/private/py_info.bzl

@@ -56,6 +56,11 @@ VenvSymlinkEntry = provider(
 An entry in `PyInfo.venv_symlinks`
 """,
     fields = {
+        "key": """
+:type: str | None
+
+Represents the value that gets used to key entries so that each package is represented only once.
+""",
         "kind": """
 :type: str
 
@@ -308,6 +313,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", "VenvSymlinkEntry", "VenvSymlinkKind")
@@ -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_venv_symlinks(ctx, semantics):
     imports = depset()
     venv_symlinks = depset()
     if VenvsSitePackages.is_enabled(ctx):
-        venv_symlinks = _get_venv_symlinks(ctx)
+        dist_info_metadata = _get_distinfo_metadata(ctx)
+        venv_symlinks = _get_venv_symlinks(
+            ctx,
+            dist_info_metadata,
+        )
     else:
         imports = collect_imports(ctx, semantics)
     return imports, venv_symlinks
 
-def _get_venv_symlinks(ctx):
+def _get_venv_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,25 @@ def _get_venv_symlinks(ctx):
     repo_runfiles_dirname = None
     dirs_with_init = {}  # dirname -> runfile path
     venv_symlinks = []
+    package = ""
+    if dist_info_metadata:
+        # in order to be able to have replacements in the venv, we have to add a
+        # third value into the venv_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("-")
+        package = normalize_name(package)
+
+        repo_runfiles_dirname = runfiles_root_path(ctx, dist_info_metadata.short_path).partition("/")[0]
+        venv_symlinks.append(VenvSymlinkEntry(
+            key = "{}.dist-info".format(package),
+            kind = VenvSymlinkKind.LIB,
+            link_to_path = paths.join(repo_runfiles_dirname, site_packages_root, dist_info_dir),
+            venv_path = dist_info_dir,
+        ))
+
     for src in ctx.files.srcs:
         if src.extension not in PYTHON_FILE_EXTENSIONS:
             continue
@@ -272,6 +308,7 @@ def _get_venv_symlinks(ctx):
             # This would be files that do not have directories and we just need to add
             # direct symlinks to them as is:
             venv_symlinks.append(VenvSymlinkEntry(
+                key = None,
                 kind = VenvSymlinkKind.LIB,
                 link_to_path = paths.join(repo_runfiles_dirname, site_packages_root, filename),
                 venv_path = filename,
@@ -293,6 +330,7 @@ def _get_venv_symlinks(ctx):
 
     for dirname in first_level_explicit_packages:
         venv_symlinks.append(VenvSymlinkEntry(
+            key = None,
             kind = VenvSymlinkKind.LIB,
             link_to_path = paths.join(repo_runfiles_dirname, site_packages_root, dirname),
             venv_path = dirname,

tests/modules/other/simple_v1/BUILD.bazel

@@ -0,0 +1,14 @@
+load("@rules_python//python:py_library.bzl", "py_library")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "simple_v1",
+    srcs = glob(["site-packages/**/*.py"]),
+    data = glob(
+        ["**/*"],
+        exclude = ["site-packages/**/*.py"],
+    ),
+    experimental_venvs_site_packages = "@rules_python//python/config_settings:venvs_site_packages",
+    imports = [package_name() + "/site-packages"],
+)

tests/modules/other/simple_v1/site-packages/simple-1.0.0.dist-info/METADATA

@@ -0,0 +1 @@
+inside is v1

tests/modules/other/simple_v1/site-packages/simple/__init__.py

@@ -0,0 +1,3 @@
+# Intentionally empty
+
+__version__ = "1.0.0"

tests/modules/other/simple_v2/BUILD.bazel

@@ -0,0 +1,14 @@
+load("@rules_python//python:py_library.bzl", "py_library")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "simple_v2",
+    srcs = glob(["site-packages/**/*.py"]),
+    data = glob(
+        ["**/*"],
+        exclude = ["site-packages/**/*.py"],
+    ),
+    experimental_venvs_site_packages = "@rules_python//python/config_settings:venvs_site_packages",
+    imports = [package_name() + "/site-packages"],
+)

tests/modules/other/simple_v2/site-packages/simple-2.0.0.dist-info/METADATA

@@ -0,0 +1 @@
+inside is v1

tests/modules/other/simple_v2/site-packages/simple/__init__.py

@@ -0,0 +1,3 @@
+# Intentionally empty
+
+__version__ = "2.0.0"

tests/venv_site_packages_libs/BUILD.bazel

@@ -1,6 +1,20 @@
+load("//python:py_library.bzl", "py_library")
 load("//tests/support:py_reconfig.bzl", "py_reconfig_test")
 load("//tests/support:support.bzl", "SUPPORTS_BOOTSTRAP_SCRIPT")
 
+py_library(
+    name = "user_lib",
+    deps = ["@other//simple_v1"],
+)
+
+py_library(
+    name = "closer_lib",
+    deps = [
+        ":user_lib",
+        "@other//simple_v2",
+    ],
+)
+
 py_reconfig_test(
     name = "venvs_site_packages_libs_test",
     srcs = ["bin.py"],
@@ -9,6 +23,7 @@ py_reconfig_test(
     target_compatible_with = SUPPORTS_BOOTSTRAP_SCRIPT,
     venvs_site_packages = "yes",
     deps = [
+        ":closer_lib",
         "//tests/venv_site_packages_libs/nspkg_alpha",
         "//tests/venv_site_packages_libs/nspkg_beta",
         "@other//nspkg_delta",

tests/venv_site_packages_libs/bin.py

@@ -2,6 +2,7 @@
 import os
 import sys
 import unittest
+from pathlib import Path
 
 
 class VenvSitePackagesLibraryTest(unittest.TestCase):
@@ -27,6 +28,24 @@ def test_imported_from_venv(self):
         self.assert_imported_from_venv("nspkg.subnspkg.gamma")
         self.assert_imported_from_venv("nspkg.subnspkg.delta")
         self.assert_imported_from_venv("single_file")
+        self.assert_imported_from_venv("simple")
+
+    def test_distinfo_is_overriden(self):
+        self.assert_imported_from_venv("simple")
+        module = importlib.import_module("simple")
+        module_path = Path(module.__file__)
+        site_packages = module_path.parent.parent
+
+        dist_info_dirs = [p.name for p in site_packages.glob("*.dist-info")]
+
+        self.assertEqual(
+            "2.0.0",
+            module.__version__,
+        )
+        self.assertEqual(
+            ["simple-2.0.0.dist-info"],
+            dist_info_dirs,
+        )
 
 
 if __name__ == "__main__":