refactor: optimize venv creation for nvidia and pkgutil style namespace packages

.bazelrc

@@ -18,6 +18,8 @@ build --//python/config_settings:incompatible_default_to_explicit_init_py=True
 
 # Ensure ongoing compatibility with this flag.
 common --incompatible_disallow_struct_provider_syntax
+# Makes Bazel 7 act more like Bazel 8
+common --incompatible_use_plus_in_repo_names
 
 # Windows makes use of runfiles for some rules
 build --enable_runfiles

AGENTS.md

@@ -27,6 +27,68 @@ end of the documentation text.
 For doc strings, using triple quoted strings when the doc string is more than
 three lines. Do not use a trailing backslack (`\`) for the opening triple-quote.
 
+### Starlark Code
+
+Starlark does not support recursion. Use iterative algorithms instead.
+
+Starlark does not support `while` loops. Use `for` loop with an appropriately
+sized iterable instead.
+
+#### Starlark testing
+
+For Starlark tests:
+
+* Use `rules_testing`, not `bazel_skylib`.
+* See https://rules-testing.readthedocs.io/en/latest/analysis_tests.html for
+  examples on using rules_testing.
+* See `tests/builders/builders_tests.bzl` for an example of using it in
+  this project.
+
+A test is defined in two parts:
+  * A setup function, e.g. `def _test_foo(name)`. This defines targets
+    and calls `analysis_test`.
+  * An implementation function, e.g. `def _test_foo_impl(env, target)`. This
+    contains asserts.
+
+Example:
+
+```
+# File: foo_tests.bzl
+
+load("@rules_testing//lib:analysis_test.bzl", "analysis_test")
+load("@rules_testing//lib:test_suite.bzl", "test_suite")
+
+_tests = []
+
+def _test_foo(name):
+    foo_library(
+        name = name + "_subject",
+    )
+    analysis_test(
+        name = name,
+        impl = _test_foo_impl,
+        target = name + "_subject",
+    )
+_tests.append(_test_foo)
+
+def _test_foo_impl(env, target):
+    env.expect.that_whatever(target[SomeInfo].whatever).equals(expected)
+
+def foo_test_suite(name):
+    test_suite(name=name, tests=_tests)
+```
+
+
+#### Repository rules
+
+The function argument `rctx` is a hint that the function is a repository rule,
+or used by a repository rule.
+
+The function argument `mrctx` is a hint that the function can be used by a
+repository rule or module extension.
+
+The `repository_ctx` API docs are at: https://bazel.build/rules/lib/builtins/repository_ctx
+
 ### bzl_library targets for bzl source files
 
 * A `bzl_library` target should be defined for every `.bzl` file outside

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

MODULE.bazel

@@ -238,6 +238,8 @@ use_repo(
     "buildkite_config",
     "implicit_namespace_ns_sub1",
     "implicit_namespace_ns_sub2",
+    "pkgutil_nspkg1",
+    "pkgutil_nspkg2",
     "rules_python_runtime_env_tc_info",
     "somepkg_with_build_files",
     "whl_with_build_files",

python/private/BUILD.bazel

@@ -675,6 +675,7 @@ bzl_library(
         "//:__subpackages__",
     ],
     deps = [
+        ":py_internal_bzl",
         "@bazel_skylib//lib:types",
     ],
 )

python/private/internal_dev_deps.bzl

@@ -67,6 +67,30 @@ def _internal_dev_deps_impl(mctx):
         enable_implicit_namespace_pkgs = False,
     )
 
+    whl_from_dir_repo(
+        name = "pkgutil_nspkg1_whl",
+        root = "//tests/repos/pkgutil_nspkg1:BUILD.bazel",
+        output = "pkgutil_nspkg1-1.0-any-none-any.whl",
+    )
+    whl_library(
+        name = "pkgutil_nspkg1",
+        whl_file = "@pkgutil_nspkg1_whl//:pkgutil_nspkg1-1.0-any-none-any.whl",
+        requirement = "pkgutil_nspkg1",
+        enable_implicit_namespace_pkgs = False,
+    )
+
+    whl_from_dir_repo(
+        name = "pkgutil_nspkg2_whl",
+        root = "//tests/repos/pkgutil_nspkg2:BUILD.bazel",
+        output = "pkgutil_nspkg2-1.0-any-none-any.whl",
+    )
+    whl_library(
+        name = "pkgutil_nspkg2",
+        whl_file = "@pkgutil_nspkg2_whl//:pkgutil_nspkg2-1.0-any-none-any.whl",
+        requirement = "pkgutil_nspkg2",
+        enable_implicit_namespace_pkgs = False,
+    )
+
 internal_dev_deps = module_extension(
     implementation = _internal_dev_deps_impl,
     doc = "This extension creates internal rules_python dev dependencies.",

python/private/py_library.bzl

@@ -99,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(
@@ -251,6 +265,7 @@ def _get_imports_and_venv_symlinks(ctx, semantics):
             package,
             version_str,
             site_packages_root = imports[0],
+            namespace_package_files = ctx.files.namespace_package_files,
         )
     else:
         imports = collect_imports(ctx, semantics)

python/private/pypi/BUILD.bazel

@@ -431,16 +431,16 @@ bzl_library(
         ":attrs_bzl",
         ":deps_bzl",
         ":generate_whl_library_build_bazel_bzl",
+        ":parse_whl_name_bzl",
         ":patch_whl_bzl",
-        ":pep508_requirement_bzl",
         ":pypi_repo_utils_bzl",
         ":whl_metadata_bzl",
         ":whl_target_platforms_bzl",
         "//python/private:auth_bzl",
-        "//python/private:bzlmod_enabled_bzl",
         "//python/private:envsubst_bzl",
         "//python/private:is_standalone_interpreter_bzl",
         "//python/private:repo_utils_bzl",
+        "//python/private:util_bzl",
         "@rules_python_internal//:rules_python_config_bzl",
     ],
 )

python/private/pypi/pypi_repo_utils.bzl

@@ -16,6 +16,7 @@
 
 load("@bazel_skylib//lib:types.bzl", "types")
 load("//python/private:repo_utils.bzl", "repo_utils")
+load("//python/private:util.bzl", "is_importable_name")
 
 def _get_python_interpreter_attr(mrctx, *, python_interpreter = None):
     """A helper function for getting the `python_interpreter` attribute or it's default
@@ -161,9 +162,44 @@ def _execute_checked_stdout(mrctx, *, python, srcs, **kwargs):
         **_execute_prep(mrctx, python = python, srcs = srcs, **kwargs)
     )
 
+def _find_namespace_package_files(rctx, install_dir):
+    """Finds all `__init__.py` files that belong to namespace packages.
+
+    A `__init__.py` file belongs to a namespace package if it contains `__path__ =`,
+    `pkgutil`, and `extend_path(`.
+
+    Args:
+        rctx (repository_ctx): The repository context.
+        install_dir (path): The path to the install directory.
+
+    Returns:
+        list[str]: A list of relative paths to `__init__.py` files that belong
+            to namespace packages.
+    """
+
+    repo_root = str(rctx.path(".")) + "/"
+    namespace_package_files = []
+    for top_level_dir in install_dir.readdir():
+        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
+        content = rctx.read(init_py)
+
+        # Look for code resembling the pkgutil namespace setup code:
+        # __path__ = __import__("pkgutil").extend_path(__path__, __name__)
+        if ("__path__ =" in content and
+            "pkgutil" in content and
+            "extend_path(" in content):
+            namespace_package_files.append(str(init_py).removeprefix(repo_root))
+
+    return namespace_package_files
+
 pypi_repo_utils = struct(
     construct_pythonpath = _construct_pypath,
     execute_checked = _execute_checked,
     execute_checked_stdout = _execute_checked_stdout,
+    find_namespace_package_files = _find_namespace_package_files,
     resolve_python_interpreter = _resolve_python_interpreter,
 )

python/private/pypi/whl_library.bzl

@@ -384,11 +384,13 @@ def _whl_library_impl(rctx):
             supports_whl_extraction = rp_config.supports_whl_extraction,
         )
 
+        install_dir_path = whl_path.dirname.get_child("site-packages")
         metadata = whl_metadata(
-            install_dir = whl_path.dirname.get_child("site-packages"),
+            install_dir = install_dir_path,
             read_fn = rctx.read,
             logger = logger,
         )
+        namespace_package_files = pypi_repo_utils.find_namespace_package_files(rctx, install_dir_path)
 
         # NOTE @aignas 2024-06-22: this has to live on until we stop supporting
         # passing `twine` as a `:pkg` library via the `WORKSPACE` builds.
@@ -432,6 +434,7 @@ def _whl_library_impl(rctx):
             data_exclude = rctx.attr.pip_data_exclude,
             group_deps = rctx.attr.group_deps,
             group_name = rctx.attr.group_name,
+            namespace_package_files = namespace_package_files,
         )
     else:
         target_platforms = rctx.attr.experimental_target_platforms or []
@@ -491,6 +494,8 @@ def _whl_library_impl(rctx):
             )
             entry_points[entry_point_without_py] = entry_point_script_name
 
+        namespace_package_files = pypi_repo_utils.find_namespace_package_files(rctx, rctx.path("site-packages"))
+
         build_file_contents = generate_whl_library_build_bazel(
             name = whl_path.basename,
             sdist_filename = sdist_filename,
@@ -509,6 +514,7 @@ def _whl_library_impl(rctx):
                 "pypi_name={}".format(metadata["name"]),
                 "pypi_version={}".format(metadata["version"]),
             ],
+            namespace_package_files = namespace_package_files,
         )
 
     # Delete these in case the wheel had them. They generally don't cause

python/private/pypi/whl_library_targets.bzl

@@ -123,6 +123,7 @@ def whl_library_targets(
         entry_points = {},
         native = native,
         enable_implicit_namespace_pkgs = False,
+        namespace_package_files = [],
         rules = struct(
             copy_file = copy_file,
             py_binary = py_binary,
@@ -169,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])
@@ -365,14 +368,16 @@ def whl_library_targets(
         )
 
         if not enable_implicit_namespace_pkgs:
-            srcs = srcs + select({
+            generated_namespace_package_files = select({
                 Label("//python/config_settings:is_venvs_site_packages"): [],
                 "//conditions:default": rules.create_inits(
                     srcs = srcs + data + pyi_srcs,
                     ignored_dirnames = [],  # If you need to ignore certain folders, you can patch rules_python here to do so.
                     root = "site-packages",
                 ),
             })
+            namespace_package_files += generated_namespace_package_files
+            srcs = srcs + generated_namespace_package_files
 
         rules.py_library(
             name = py_library_label,
@@ -391,6 +396,7 @@ def whl_library_targets(
             tags = tags,
             visibility = impl_vis,
             experimental_venvs_site_packages = Label("@rules_python//python/config_settings:venvs_site_packages"),
+            namespace_package_files = namespace_package_files,
         )
 
 def _config_settings(dependencies_by_platform, dependencies_with_markers, rules, native = native, **kwargs):

python/private/util.bzl

@@ -15,6 +15,7 @@
 """Functionality shared by multiple pieces of code."""
 
 load("@bazel_skylib//lib:types.bzl", "types")
+load("//python/private:py_internal.bzl", "py_internal")
 
 def copy_propagating_kwargs(from_kwargs, into_kwargs = None):
     """Copies args that must be compatible between two targets with a dependency relationship.
@@ -69,3 +70,18 @@ def add_tag(attrs, tag):
             attrs["tags"] = tags + [tag]
     else:
         attrs["tags"] = [tag]
+
+def is_importable_name(name):
+    # Requires Bazel 8+
+    if hasattr(py_internal, "regex_match"):
+        # ?U means activates unicode matching (Python allows most unicode
+        # in module names / identifiers).
+        # \w matches alphanumeric and underscore.
+        # NOTE: regex_match has an implicit ^ and $
+        return py_internal.regex_match(name, "(?U)\\w+")
+    else:
+        # Otherwise, use a rough hueristic that should catch most cases.
+        return (
+            "." not in name and
+            "-" not in name
+        )