correctly pop existing links for package overrides

Author: aignas

python/private/py_executable.bzl

@@ -680,17 +680,32 @@ def _create_venv_symlinks(ctx, venv_dir_map):
     return venv_files
 
 def _build_link_map(entries):
-    # dict[str kind, dict[str key or rel_path, tuple[str rel_path, str link_to_path]]]
+    # 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 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.key or entry.venv_path] = (entry.venv_path, entry.link_to_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
@@ -710,7 +725,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,11 +56,6 @@ 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
 
@@ -72,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

python/private/py_library.bzl

@@ -271,7 +271,7 @@ def _get_venv_symlinks(ctx, dist_info_metadata):
     repo_runfiles_dirname = None
     dirs_with_init = {}  # dirname -> runfile path
     venv_symlinks = []
-    package = ""
+    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
@@ -284,9 +284,9 @@ def _get_venv_symlinks(ctx, dist_info_metadata):
 
         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),
+            src = package,
             venv_path = dist_info_dir,
         ))
 
@@ -308,9 +308,9 @@ def _get_venv_symlinks(ctx, dist_info_metadata):
             # 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),
+                src = package,
                 venv_path = filename,
             ))
 
@@ -330,9 +330,9 @@ def _get_venv_symlinks(ctx, dist_info_metadata):
 
     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),
+            src = package,
             venv_path = dirname,
         ))
     return venv_symlinks

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

@@ -1,3 +1 @@
-# Intentionally empty
-
 __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/site-packages/simple/__init__.py

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

tests/venv_site_packages_libs/bin.py

@@ -33,20 +33,26 @@ def test_imported_from_venv(self):
     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__,
         )
+        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__":
     unittest.main()