basic impl to optimize namespace packages

Author: rickeylev

python/private/venv_runfiles.bzl

@@ -13,6 +13,7 @@ load(
     "VenvSymlinkEntry",
     "VenvSymlinkKind",
 )
+load(":py_internal.bzl", "py_internal")
 
 def create_venv_app_files(ctx, deps, venv_dir_map):
     """Creates the tree of app-specific files for a venv for a binary.
@@ -253,6 +254,9 @@ def get_venv_symlinks(ctx, files, package, version_str, site_packages_root):
     # List of (File, str venv_path) tuples
     files_left_to_link = []
 
+    # dict[str dirname, bool is_namespace_package]
+    namespace_package_dirs = {}
+
     # We want to minimize the number of files symlinked. Ideally, only the
     # top-level directories are symlinked. Unfortunately, shared libraries
     # complicate matters: if a shared library's directory is linked, then the
@@ -293,6 +297,29 @@ def get_venv_symlinks(ctx, files, package, version_str, site_packages_root):
         else:
             files_left_to_link.append((src, venv_path))
 
+        top_level_dirname, _, tail = venv_path.partition("/")
+        if (
+            # If it's already not directly linkable, nothing to do
+            not cannot_be_linked_directly.get(top_level_dirname, False) and
+            # If its already known to be non-implicit namespace, then skip
+            namespace_package_dirs.get(top_level_dirname, True) and
+            # It must be an importable name to be an implicit namespace package
+            py_internal.regex_match(top_level_dirname, "(?U)\\w+")
+        ):
+            namespace_package_dirs.setdefault(top_level_dirname, True)
+
+            # Looking for `__init__.` isn't 100% correct, as it'll match e.g.
+            # `__init__.pyi`, but it's close enough.
+            if "/" not in tail and tail.startswith("__init__."):
+                namespace_package_dirs[top_level_dirname] = False
+
+    # We treat namespace packages as a hint that other distributions may
+    # install into the same directory. As such, we avoid linking them directly
+    # to avoid conflict merging later.
+    for dirname, is_namespace_package in namespace_package_dirs.items():
+        if is_namespace_package:
+            cannot_be_linked_directly[dirname] = True
+
     # At this point, venv_symlinks has entries for the shared libraries
     # and cannot_be_linked_directly has the directories that cannot be
     # directly linked. Next, we loop over the remaining files and group

tests/venv_site_packages_libs/app_files_building/app_files_building_tests.bzl

@@ -218,6 +218,7 @@ def _test_optimized_grouping_single_toplevel(name):
     empty_files(
         name = name + "_files",
         paths = [
+            "site-packages/pkg2/__init__.py",
             "site-packages/pkg2/a.txt",
             "site-packages/pkg2/b_mod.so",
         ],
@@ -247,6 +248,7 @@ def _test_optimized_grouping_single_toplevel_impl(env, target):
             "pkg2",
             link_to_path = rr + "pkg2",
             files = [
+                "tests/venv_site_packages_libs/app_files_building/site-packages/pkg2/__init__.py",
                 "tests/venv_site_packages_libs/app_files_building/site-packages/pkg2/a.txt",
                 "tests/venv_site_packages_libs/app_files_building/site-packages/pkg2/b_mod.so",
             ],
@@ -263,6 +265,68 @@ def _test_optimized_grouping_single_toplevel_impl(env, target):
     # The point of the optimization is to avoid having to merge conflicts.
     env.expect.that_collection(conflicts).contains_exactly([])
 
+def _test_optimized_grouping_implicit_namespace_packages(name):
+    empty_files(
+        name = name + "_files",
+        paths = [
+            "site-packages/namespace/part1/foo.py",
+            "site-packages/namespace/part2/bar.py",
+            "site-packages/namespace-1.0.dist-info/METADATA",
+        ],
+    )
+    analysis_test(
+        name = name,
+        impl = _test_optimized_grouping_implicit_namespace_packages_impl,
+        target = name + "_files",
+    )
+
+_tests.append(_test_optimized_grouping_implicit_namespace_packages)
+
+def _test_optimized_grouping_implicit_namespace_packages_impl(env, target):
+    test_ctx = _ctx(workspace_name = env.ctx.workspace_name)
+    entries = get_venv_symlinks(
+        test_ctx,
+        target.files.to_list(),
+        package = "pkg3",
+        version_str = "1.0",
+        site_packages_root = env.ctx.label.package + "/site-packages",
+    )
+    actual = _venv_symlinks_from_entries(entries)
+
+    rr = "{}/{}/site-packages/".format(test_ctx.workspace_name, env.ctx.label.package)
+    expected = [
+        _venv_symlink(
+            "namespace/part1",
+            link_to_path = rr + "namespace/part1",
+            files = [
+                "tests/venv_site_packages_libs/app_files_building/site-packages/namespace/part1/foo.py",
+            ],
+        ),
+        _venv_symlink(
+            "namespace/part2",
+            link_to_path = rr + "namespace/part2",
+            files = [
+                "tests/venv_site_packages_libs/app_files_building/site-packages/namespace/part2/bar.py",
+            ],
+        ),
+        _venv_symlink(
+            "namespace-1.0.dist-info",
+            link_to_path = rr + "namespace-1.0.dist-info",
+            files = [
+                "tests/venv_site_packages_libs/app_files_building/site-packages/namespace-1.0.dist-info/METADATA",
+            ],
+        ),
+    ]
+    expected = sorted(expected, key = lambda e: (e.link_to_path, e.venv_path))
+    env.expect.that_collection(
+        actual,
+    ).contains_exactly(expected)
+
+    _, conflicts = build_link_map(test_ctx, entries, return_conflicts = True)
+
+    # The point of the optimization is to avoid having to merge conflicts.
+    env.expect.that_collection(conflicts).contains_exactly([])
+
 def _test_package_version_filtering(name):
     analysis_test(
         name = name,