refactor: reimplement writing namespace pkgs in Starlark (#2882)

With this PR I would like to facilitate the implementation of the venv
layouts because we can in theory take the `srcs` and the `data` within
the `py_library` and then use the `expand_template` to write the extra
Python files if the namespace_pkgs flag is enabled.

The old Python code has been removed and the extra generated files are
written out with `bazel_skylib` `copy_file`.

The implicit `namespace_pkg` init files are included to `py_library`
if the `site-packages` config flag is set to false and I think this
may help with continuing the implementation, but it currently is still
not working as expected (see comment).

Work towards #2156

Author: aignas

python/config_settings/BUILD.bazel

@@ -217,6 +217,15 @@ string_flag(
     visibility = ["//visibility:public"],
 )
 
+config_setting(
+    name = "is_venvs_site_packages",
+    flag_values = {
+        ":venvs_site_packages": VenvsSitePackages.YES,
+    },
+    # NOTE: Only public because it is used in whl_library repos.
+    visibility = ["//visibility:public"],
+)
+
 define_pypi_internal_flags(
     name = "define_pypi_internal_flags",
 )

python/private/pypi/BUILD.bazel

@@ -18,6 +18,11 @@ package(default_visibility = ["//:__subpackages__"])
 
 licenses(["notice"])
 
+exports_files(
+    srcs = ["namespace_pkg_tmpl.py"],
+    visibility = ["//visibility:public"],
+)
+
 filegroup(
     name = "distribution",
     srcs = glob(

python/private/pypi/namespace_pkg_tmpl.py

@@ -0,0 +1,2 @@
+# __path__ manipulation added by bazel-contrib/rules_python to support namespace pkgs.
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

python/private/pypi/namespace_pkgs.bzl

@@ -0,0 +1,83 @@
+"""Utilities to get where we should write namespace pkg paths."""
+
+load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
+
+_ext = struct(
+    py = ".py",
+    pyd = ".pyd",
+    so = ".so",
+    pyc = ".pyc",
+)
+
+_TEMPLATE = Label("//python/private/pypi:namespace_pkg_tmpl.py")
+
+def _add_all(dirname, dirs):
+    dir_path = "."
+    for dir_name in dirname.split("/"):
+        dir_path = "{}/{}".format(dir_path, dir_name)
+        dirs[dir_path[2:]] = None
+
+def get_files(*, srcs, ignored_dirnames = [], root = None):
+    """Get the list of filenames to write the namespace pkg files.
+
+    Args:
+        srcs: {type}`src` a list of files to be passed to {bzl:obj}`py_library`
+            as `srcs` and `data`. This is usually a result of a {obj}`glob`.
+        ignored_dirnames: {type}`str` a list of patterns to ignore.
+        root: {type}`str` the prefix to use as the root.
+
+    Returns:
+        {type}`src` a list of paths to write the namespace pkg `__init__.py` file.
+    """
+    dirs = {}
+    ignored = {i: None for i in ignored_dirnames}
+
+    if root:
+        _add_all(root, ignored)
+
+    for file in srcs:
+        dirname, _, filename = file.rpartition("/")
+
+        if filename == "__init__.py":
+            ignored[dirname] = None
+            dirname, _, _ = dirname.rpartition("/")
+        elif filename.endswith(_ext.py):
+            pass
+        elif filename.endswith(_ext.pyc):
+            pass
+        elif filename.endswith(_ext.pyd):
+            pass
+        elif filename.endswith(_ext.so):
+            pass
+        else:
+            continue
+
+        if dirname in dirs or not dirname:
+            continue
+
+        _add_all(dirname, dirs)
+
+    return sorted([d for d in dirs if d not in ignored])
+
+def create_inits(**kwargs):
+    """Create init files and return the list to be included `py_library` srcs.
+
+    Args:
+        **kwargs: passed to {obj}`get_files`.
+
+    Returns:
+        {type}`list[str]` to be included as part of `py_library`.
+    """
+    srcs = []
+    for out in get_files(**kwargs):
+        src = "{}/__init__.py".format(out)
+        srcs.append(srcs)
+
+        copy_file(
+            name = "_cp_{}_namespace".format(out),
+            src = _TEMPLATE,
+            out = src,
+            **kwargs
+        )
+
+    return srcs

python/private/pypi/whl_installer/arguments.py

@@ -57,11 +57,6 @@ def parser(**kwargs: Any) -> argparse.ArgumentParser:
         action="store",
         help="Additional data exclusion parameters to add to the pip packages BUILD file.",
     )
-    parser.add_argument(
-        "--enable_implicit_namespace_pkgs",
-        action="store_true",
-        help="Disables conversion of implicit namespace packages into pkg-util style packages.",
-    )
     parser.add_argument(
         "--environment",
         action="store",

python/private/pypi/whl_installer/wheel_installer.py

@@ -27,7 +27,7 @@
 
 from pip._vendor.packaging.utils import canonicalize_name
 
-from python.private.pypi.whl_installer import arguments, namespace_pkgs, wheel
+from python.private.pypi.whl_installer import arguments, wheel
 
 
 def _configure_reproducible_wheels() -> None:
@@ -77,35 +77,10 @@ def _parse_requirement_for_extra(
     return None, None
 
 
-def _setup_namespace_pkg_compatibility(wheel_dir: str) -> None:
-    """Converts native namespace packages to pkgutil-style packages
-
-    Namespace packages can be created in one of three ways. They are detailed here:
-    https://packaging.python.org/guides/packaging-namespace-packages/#creating-a-namespace-package
-
-    'pkgutil-style namespace packages' (2) and 'pkg_resources-style namespace packages' (3) works in Bazel, but
-    'native namespace packages' (1) do not.
-
-    We ensure compatibility with Bazel of method 1 by converting them into method 2.
-
-    Args:
-        wheel_dir: the directory of the wheel to convert
-    """
-
-    namespace_pkg_dirs = namespace_pkgs.implicit_namespace_packages(
-        wheel_dir,
-        ignored_dirnames=["%s/bin" % wheel_dir],
-    )
-
-    for ns_pkg_dir in namespace_pkg_dirs:
-        namespace_pkgs.add_pkgutil_style_namespace_pkg_init(ns_pkg_dir)
-
-
 def _extract_wheel(
     wheel_file: str,
     extras: Dict[str, Set[str]],
     enable_pipstar: bool,
-    enable_implicit_namespace_pkgs: bool,
     platforms: List[wheel.Platform],
     installation_dir: Path = Path("."),
 ) -> None:
@@ -116,15 +91,11 @@ def _extract_wheel(
         installation_dir: the destination directory for installation of the wheel.
         extras: a list of extras to add as dependencies for the installed wheel
         enable_pipstar: if true, turns off certain operations.
-        enable_implicit_namespace_pkgs: if true, disables conversion of implicit namespace packages and will unzip as-is
     """
 
     whl = wheel.Wheel(wheel_file)
     whl.unzip(installation_dir)
 
-    if not enable_implicit_namespace_pkgs:
-        _setup_namespace_pkg_compatibility(installation_dir)
-
     metadata = {
         "entry_points": [
             {
@@ -168,7 +139,6 @@ def main() -> None:
             wheel_file=whl,
             extras=extras,
             enable_pipstar=args.enable_pipstar,
-            enable_implicit_namespace_pkgs=args.enable_implicit_namespace_pkgs,
             platforms=arguments.get_platforms(args),
         )
         return

python/private/pypi/whl_library.bzl

@@ -173,9 +173,6 @@ def _parse_optional_attrs(rctx, args, extra_pip_args = None):
             json.encode(struct(arg = rctx.attr.pip_data_exclude)),
         ]
 
-    if rctx.attr.enable_implicit_namespace_pkgs:
-        args.append("--enable_implicit_namespace_pkgs")
-
     env = {}
     if rctx.attr.environment != None:
         for key, value in rctx.attr.environment.items():
@@ -389,6 +386,8 @@ def _whl_library_impl(rctx):
             metadata_name = metadata.name,
             metadata_version = metadata.version,
             requires_dist = metadata.requires_dist,
+            # TODO @aignas 2025-05-17: maybe have a build flag for this instead
+            enable_implicit_namespace_pkgs = rctx.attr.enable_implicit_namespace_pkgs,
             # TODO @aignas 2025-04-14: load through the hub:
             annotation = None if not rctx.attr.annotation else struct(**json.decode(rctx.read(rctx.attr.annotation))),
             data_exclude = rctx.attr.pip_data_exclude,
@@ -457,6 +456,8 @@ def _whl_library_impl(rctx):
             name = whl_path.basename,
             dep_template = rctx.attr.dep_template or "@{}{{name}}//:{{target}}".format(rctx.attr.repo_prefix),
             entry_points = entry_points,
+            # TODO @aignas 2025-05-17: maybe have a build flag for this instead
+            enable_implicit_namespace_pkgs = rctx.attr.enable_implicit_namespace_pkgs,
             # TODO @aignas 2025-04-14: load through the hub:
             dependencies = metadata["deps"],
             dependencies_by_platform = metadata["deps_by_platform"],
@@ -580,7 +581,6 @@ attr makes `extra_pip_args` and `download_only` ignored.""",
             Label("//python/private/pypi/whl_installer:wheel.py"),
             Label("//python/private/pypi/whl_installer:wheel_installer.py"),
             Label("//python/private/pypi/whl_installer:arguments.py"),
-            Label("//python/private/pypi/whl_installer:namespace_pkgs.py"),
         ] + record_files.values(),
     ),
     "_rule_name": attr.string(default = "whl_library"),

python/private/pypi/whl_library_targets.bzl

@@ -30,6 +30,7 @@ load(
     "WHEEL_FILE_IMPL_LABEL",
     "WHEEL_FILE_PUBLIC_LABEL",
 )
+load(":namespace_pkgs.bzl", "create_inits")
 load(":pep508_deps.bzl", "deps")
 
 def whl_library_targets_from_requires(
@@ -113,6 +114,7 @@ def whl_library_targets(
         copy_executables = {},
         entry_points = {},
         native = native,
+        enable_implicit_namespace_pkgs = False,
         rules = struct(
             copy_file = copy_file,
             py_binary = py_binary,
@@ -153,6 +155,8 @@ def whl_library_targets(
         data: {type}`list[str]` A list of labels to include as part of the `data` attribute in `py_library`.
         entry_points: {type}`dict[str, str]` The mapping between the script
             name and the python file to use. DEPRECATED.
+        enable_implicit_namespace_pkgs: {type}`boolean` generate __init__.py
+            files for namespace pkgs.
         native: {type}`native` The native struct for overriding in tests.
         rules: {type}`struct` A struct with references to rules for creating targets.
     """
@@ -293,6 +297,14 @@ def whl_library_targets(
         )
 
     if hasattr(rules, "py_library"):
+        srcs = native.glob(
+            ["site-packages/**/*.py"],
+            exclude = srcs_exclude,
+            # Empty sources are allowed to support wheels that don't have any
+            # pure-Python code, e.g. pymssql, which is written in Cython.
+            allow_empty = True,
+        )
+
         # NOTE: pyi files should probably be excluded because they're carried
         # by the pyi_srcs attribute. However, historical behavior included
         # them in data and some tools currently rely on that.
@@ -309,23 +321,31 @@ def whl_library_targets(
             if item not in _data_exclude:
                 _data_exclude.append(item)
 
+        data = data + native.glob(
+            ["site-packages/**/*"],
+            exclude = _data_exclude,
+        )
+
+        pyi_srcs = native.glob(
+            ["site-packages/**/*.pyi"],
+            allow_empty = True,
+        )
+
+        if enable_implicit_namespace_pkgs:
+            srcs = srcs + getattr(native, "select", select)({
+                Label("//python/config_settings:is_venvs_site_packages"): [],
+                "//conditions:default": create_inits(
+                    srcs = srcs + data + pyi_srcs,
+                    ignore_dirnames = [],  # If you need to ignore certain folders, you can patch rules_python here to do so.
+                    root = "site-packages",
+                ),
+            })
+
         rules.py_library(
             name = py_library_label,
-            srcs = native.glob(
-                ["site-packages/**/*.py"],
-                exclude = srcs_exclude,
-                # Empty sources are allowed to support wheels that don't have any
-                # pure-Python code, e.g. pymssql, which is written in Cython.
-                allow_empty = True,
-            ),
-            pyi_srcs = native.glob(
-                ["site-packages/**/*.pyi"],
-                allow_empty = True,
-            ),
-            data = data + native.glob(
-                ["site-packages/**/*"],
-                exclude = _data_exclude,
-            ),
+            srcs = srcs,
+            pyi_srcs = pyi_srcs,
+            data = data,
             # This makes this directory a top-level in the python import
             # search path for anything that depends on this.
             imports = ["site-packages"],

tests/pypi/namespace_pkgs/BUILD.bazel

@@ -0,0 +1,5 @@
+load(":namespace_pkgs_tests.bzl", "namespace_pkgs_test_suite")
+
+namespace_pkgs_test_suite(
+    name = "namespace_pkgs_tests",
+)