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,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.
-        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[VenvSymlinkEntry]` instances 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

@@ -682,15 +682,30 @@ def _create_venv_symlinks(ctx, venv_dir_map):
 def _build_link_map(entries):
     # dict[str kind, dict[str rel_path, str link_to_path]]
     link_map = {}
+
+    # Here we store venv paths by package
+    # dict[str package, dict[str kind, list[str venv_path]]]
+    pkg_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
+
+        # If we detect that we are adding a dist-info for an already existing package
+        # we need to pop all of the previous symlinks from the link_map
+        if entry.venv_path.endswith(".dist-info") and entry.src in pkg_map:
+            # dist-info will come always first
+            for kind, dir_paths in pkg_map.pop(entry.src).items():
+                for dir_path in dir_paths:
+                    link_map[kind].pop(dir_path)
+
+        pkg_venv_paths = pkg_map.setdefault(entry.src, {}).setdefault(entry.kind, [])
+        pkg_venv_paths.append(entry.venv_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.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

python/private/py_info.bzl

@@ -67,6 +67,11 @@ the venv to create the path under.
 
 A runfiles-root relative path that `venv_path` will symlink to. If `None`,
 it means to not create a symlink.
+""",
+        "src": """
+:type: str | None
+
+Represents the PyPI package that the code originates from.
 """,
         "venv_path": """
 :type: str
@@ -298,17 +303,6 @@ This field is currently unused in Bazel and may go away in the future.
 
 A depset with `topological` ordering.
 
-
-Tuples of `(runfiles_path, site_packages_path)`. 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`
-  name.
-* `site_packages_path` is a path relative to the site-packages directory of
-  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`).
-
 :::{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,16 +271,41 @@ def _get_venv_symlinks(ctx):
     repo_runfiles_dirname = None
     dirs_with_init = {}  # dirname -> runfile path
     venv_symlinks = []
-    for src in ctx.files.srcs:
-        if src.extension not in PYTHON_FILE_EXTENSIONS:
-            continue
+    package = None
+    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(
+            kind = VenvSymlinkKind.LIB,
+            link_to_path = paths.join(repo_runfiles_dirname, site_packages_root, dist_info_dir),
+            src = package,
+            venv_path = dist_info_dir,
+        ))
+
+    for src in ctx.files.srcs + ctx.files.data:
         path = _repo_relative_short_path(src.short_path)
         if not path.startswith(site_packages_root):
             continue
         path = path.removeprefix(site_packages_root)
         dir_name, _, filename = path.rpartition("/")
 
-        if dir_name and filename.startswith("__init__."):
+        if src.extension not in PYTHON_FILE_EXTENSIONS:
+            if dir_name.endswith(".dist-info"):
+                # we have already handled the stuff
+                pass
+            elif dir_name:
+                # TODO @aignas 2025-05-30: is this the right way?
+                dirs_with_init[dir_name] = None
+                repo_runfiles_dirname = runfiles_root_path(ctx, src.short_path).partition("/")[0]
+        elif dir_name and filename.startswith("__init__."):
             dirs_with_init[dir_name] = None
             repo_runfiles_dirname = runfiles_root_path(ctx, src.short_path).partition("/")[0]
         elif not dir_name:
@@ -274,6 +316,7 @@ def _get_venv_symlinks(ctx):
             venv_symlinks.append(VenvSymlinkEntry(
                 kind = VenvSymlinkKind.LIB,
                 link_to_path = paths.join(repo_runfiles_dirname, site_packages_root, filename),
+                src = package,
                 venv_path = filename,
             ))
 
@@ -295,6 +338,7 @@ def _get_venv_symlinks(ctx):
         venv_symlinks.append(VenvSymlinkEntry(
             kind = VenvSymlinkKind.LIB,
             link_to_path = paths.join(repo_runfiles_dirname, site_packages_root, dirname),
+            src = package,
             venv_path = dirname,
         ))
     return venv_symlinks

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 @@
+__version__ = "1.0.0"

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

@@ -0,0 +1 @@
+# Intentionally empty

tests/modules/other/simple_v2/BUILD.bazel

@@ -0,0 +1,15 @@
+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"],
+    pyi_srcs = glob(["**/*.pyi"]),
+)

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.libs/data.txt

@@ -0,0 +1 @@
+Intentionally empty

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

@@ -0,0 +1 @@
+__version__ = "2.0.0"

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

@@ -0,0 +1 @@
+# Intentionally empty

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

@@ -1,7 +1,7 @@
 import importlib
-import os
 import sys
 import unittest
+from pathlib import Path
 
 
 class VenvSitePackagesLibraryTest(unittest.TestCase):
@@ -27,6 +27,51 @@ 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_pyi_is_included(self):
+        self.assert_imported_from_venv("simple")
+        module = importlib.import_module("simple")
+        module_path = Path(module.__file__)
+
+        # this has been not included through data but through `pyi_srcs`
+        pyi_files = [p.name for p in module_path.parent.glob("*.pyi")]
+        self.assertIn("__init__.pyi", pyi_files)
+
+    def test_data_is_included(self):
+        self.assert_imported_from_venv("simple")
+        module = importlib.import_module("simple")
+        module_path = Path(module.__file__)
+
+        site_packages = module_path.parent.parent
+
+        # Ensure that packages from simple v1 are not present
+        files = [p.name for p in site_packages.glob("*")]
+        self.assertIn("simple.libs", files)
+
+    def test_override_pkg(self):
+        self.assert_imported_from_venv("simple")
+        module = importlib.import_module("simple")
+        self.assertEqual(
+            "2.0.0",
+            module.__version__,
+        )
+
+    def test_dirs_from_replaced_package_are_not_present(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(
+            ["simple-2.0.0.dist-info"],
+            dist_info_dirs,
+        )
+
+        # Ensure that packages from simple v1 are not present
+        files = [p.name for p in site_packages.glob("*")]
+        self.assertNotIn("simple_extras", files)
 
 
 if __name__ == "__main__":