Merge branch 'refactor.optimize.venv.namespace.packages' into refactor.optimize.venv.pkgutil.namespace.packages

Author: rickeylev

CHANGELOG.md

@@ -118,6 +118,8 @@ END_UNRELEASED_TEMPLATE
   `RULES_PYTHON_ENABLE_PIPSTAR=1` by default. Users of `experimental_index_url` that perform
   cross-builds should add {obj}`target_platforms` to their `pip.parse` invocations, which will
   become mandatory if any cross-builds are required from the next release.
+* (py_library) Attribute {obj}`namespace_package_files` added. It is a hint for
+  optimizing venv creation.
 
 [20251031]: https://github.com/astral-sh/python-build-standalone/releases/tag/20251031
 [20251202]: https://github.com/astral-sh/python-build-standalone/releases/tag/20251202

python/private/internal_dev_deps.bzl

@@ -67,18 +67,6 @@ def _internal_dev_deps_impl(mctx):
         enable_implicit_namespace_pkgs = False,
     )
 
-    whl_from_dir_repo(
-        name = "pkgutil_namespace_whl",
-        root = "//tests/pkgutil_namespace:testdata/ns-sub2/__init__.py",
-        output = "my_pkgutil_ns-1.0-any-none-any.whl",
-    )
-    whl_library(
-        name = "pkgutil_namespace",
-        whl_file = "@pkgutil_namespace_whl//:my_pkgutil_ns-1.0-any-none-any.whl",
-        requirement = "my-pkgutil-ns",
-        enable_implicit_namespace_pkgs = False,
-    )
-
     whl_from_dir_repo(
         name = "pkgutil_nspkg1_whl",
         root = "//tests/repos/pkgutil_nspkg1:BUILD.bazel",

python/private/py_library.bzl

@@ -59,11 +59,6 @@ LIBRARY_ATTRS = dicts.add(
     PY_SRCS_ATTRS,
     IMPORTS_ATTRS,
     {
-        "namespace_package_files": lambda: attrb.LabelList(
-            allow_empty = True,
-            allow_files = True,
-            doc = "The files that make up a namespace package.",
-        ),
         "experimental_venvs_site_packages": lambda: attrb.Label(
             doc = """
 **INTERNAL ATTRIBUTE. SHOULD ONLY BE SET BY rules_python-INTERNAL CODE.**
@@ -104,6 +99,20 @@ The topological order has been removed and if 2 different versions of the same P
 package are observed, the behaviour has no guarantees except that it is deterministic
 and that only one package version will be included.
 :::
+""",
+        ),
+        "namespace_package_files": lambda: attrb.LabelList(
+            allow_empty = True,
+            allow_files = True,
+            doc = """
+Files whose directories are namespace packages.
+
+When {obj}`--venv_site_packages=yes` is set, this helps inform which directories should be
+treated as namespace packages and expect files from other targets to be contributed.
+This allows optimizing the generation of symlinks to be cheaper at analysis time.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
 """,
         ),
         "_add_srcs_to_runfiles_flag": lambda: attrb.Label(

python/private/pypi/whl_library.bzl

@@ -18,7 +18,6 @@ load("@rules_python_internal//:rules_python_config.bzl", rp_config = "config")
 load("//python/private:auth.bzl", "AUTH_ATTRS", "get_auth")
 load("//python/private:envsubst.bzl", "envsubst")
 load("//python/private:is_standalone_interpreter.bzl", "is_standalone_interpreter")
-load("//python/private:py_internal.bzl", "py_internal")
 load("//python/private:repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")
 load("//python/private:util.bzl", "is_importable_name")
 load(":attrs.bzl", "ATTRS", "use_isolated")
@@ -590,18 +589,13 @@ def _find_namespace_package_files(rctx, install_dir):
     """
 
     repo_root = str(rctx.path(".")) + "/"
-    if "nvidia" in repo_root:
-        print("============== nvidia")
     namespace_package_files = []
-    print("install_dir:", install_dir)
     for top_level_dir in install_dir.readdir():
-        print("tld:", top_level_dir)
         if not is_importable_name(top_level_dir.basename):
             continue
         init_py = top_level_dir.get_child("__init__.py")
         if not init_py.exists:
             continue
-        print("found tld init:", init_py)
         content = rctx.read(init_py)
 
         # Look for code resembling the pkgutil namespace setup code:

python/private/pypi/whl_library_targets.bzl

@@ -170,6 +170,8 @@ def whl_library_targets(
         enable_implicit_namespace_pkgs: {type}`boolean` generate __init__.py
             files for namespace pkgs.
         native: {type}`native` The native struct for overriding in tests.
+        namespace_package_files: {type}`list[str]` A list of labels of files whose
+            directories are namespace packages.
         rules: {type}`struct` A struct with references to rules for creating targets.
     """
     dependencies = sorted([normalize_name(d) for d in dependencies])

python/private/venv_runfiles.bzl

@@ -13,7 +13,6 @@ load(
     "VenvSymlinkEntry",
     "VenvSymlinkKind",
 )
-load(":py_internal.bzl", "py_internal")
 load(":util.bzl", "is_importable_name")
 
 # List of top-level package names that are known to be namespace
@@ -140,8 +139,6 @@ def build_link_map(ctx, entries, return_conflicts = False):
                 else:
                     keep_kind_link_map[entry.venv_path] = entry.link_to_path
             else:
-                for v in group:
-                    print("conflict:", v.venv_path, v.link_to_path)
                 if return_conflicts:
                     conflicts.append(group)
 
@@ -313,18 +310,14 @@ def get_venv_symlinks(
         dirname: True
         for dirname in namespace_package_dirs.keys()
     }
-    print("npf:::", ctx.label, namespace_package_files)
     for f in namespace_package_files:
         venv_path, _ = _get_file_venv_path(ctx, f, site_packages_root)
         if venv_path == None:
             continue
         ns_dir = paths.dirname(venv_path)
-        print("nsdir:", ns_dir)
         namespace_package_dirs[ns_dir] = True
         cannot_be_linked_directly[ns_dir] = True
 
-    print("all npd:", namespace_package_dirs)
-
     # dict[str path, VenvSymlinkEntry]
     # Where path is the venv path (i.e. relative to site_packages_prefix)
     venv_symlinks = {}

tests/repos/pkgutil_nspkg1/nspkg/__init__.py

@@ -1,2 +1,2 @@
 # __init__.py
-__path__ = __import__('pkgutil').extend_path(__path__, __name__)
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

tests/repos/pkgutil_nspkg2/nspkg/__init__.py

@@ -1,2 +1,2 @@
 # __init__.py
-__path__ = __import__('pkgutil').extend_path(__path__, __name__)
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

tests/venv_site_packages_libs/app_files_building/app_files_building_tests.bzl

@@ -4,7 +4,7 @@ load("@rules_testing//lib:analysis_test.bzl", "analysis_test")
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("//python:py_info.bzl", "PyInfo")
 load("//python:py_library.bzl", "py_library")
-load("//python/private:common_labels.bzl", "labels")
+load("//python/private:common_labels.bzl", "labels")  # buildifier: disable=bzl-visibility
 load("//python/private:py_info.bzl", "VenvSymlinkEntry", "VenvSymlinkKind")  # buildifier: disable=bzl-visibility
 load("//python/private:venv_runfiles.bzl", "build_link_map", "get_venv_symlinks")  # buildifier: disable=bzl-visibility
 
@@ -340,6 +340,9 @@ def _test_optimized_grouping_pkgutil_namespace_packages(name):
         paths = [
             "site-packages/pkgutilns/__init__.py",
             "site-packages/pkgutilns/foo.py",
+            # Special cases: These dirnames under site-packages are always
+            # treated as namespace packages
+            "site-packages/nvidia/whatever/w.py",
         ],
     )
     analysis_test(
@@ -353,7 +356,6 @@ _tests.append(_test_optimized_grouping_pkgutil_namespace_packages)
 def _test_optimized_grouping_pkgutil_namespace_packages_impl(env, target):
     test_ctx = _ctx(workspace_name = env.ctx.workspace_name)
     files = target.files.to_list()
-    pkgutilns_init_py = None
     ns_inits = [f for f in files if f.basename == "__init__.py"]
 
     entries = get_venv_symlinks(
@@ -382,6 +384,13 @@ def _test_optimized_grouping_pkgutil_namespace_packages_impl(env, target):
                 "tests/venv_site_packages_libs/app_files_building/site-packages/pkgutilns/foo.py",
             ],
         ),
+        _venv_symlink(
+            "nvidia/whatever",
+            link_to_path = rr + "nvidia/whatever",
+            files = [
+                "tests/venv_site_packages_libs/app_files_building/site-packages/nvidia/whatever/w.py",
+            ],
+        ),
     ]
     expected = sorted(expected, key = lambda e: (e.link_to_path, e.venv_path))
     env.expect.that_collection(