feat: add an env variable to toggle pipstar (#2855)

aignas · rickeylev · web-flow · commit a4b946bbe1b3 · 2025-05-05T05:22:38.000Z
This is a flag to start leveraging of the new code paths. The Starlark implementation has been added in 1.4 and has been reverted in the latest release candidates. The `env` variable will be a good way to roll it out more gradually and get more testing. For now we are switching only the `whl_library` internals as the `requirements.txt` files from `uv` may use `*` in `python_full_version` and `platform_version` that are not yet fully supported (#2826). Main goals for this is to start using Starlark implementation so that we don't have any hidden variables. What is more, having this in Starlark is the most maintainable long-term solution for supporting cross-platform builds. Work towards #260 --------- Co-authored-by: Richard Levasseur <richardlev@gmail.com>
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -86,6 +86,9 @@ END_UNRELEASED_TEMPLATE
   (the default), the subprocess's stdout/stderr will be logged.
 * (toolchains) Local toolchains can be activated with custom flags. See
   [Conditionally using local toolchains] docs for how to configure.
+* (pypi) `RULES_PYTHON_ENABLE_PIPSTAR` environment variable: when `1`, the Starlark
+  implementation of wheel METADATA parsing is used (which has improved multi-platform
+  build support).
 
 {#v0-0-0-removed}
 ### Removed
diff --git a/docs/environment-variables.md b/docs/environment-variables.md
@@ -60,6 +60,15 @@ The default became `1` if unspecified
 :::
 ::::
 
+::::{envvar} RULES_PYTHON_ENABLE_PIPSTAR
+
+When `1`, the rules_python Starlark implementation of the pypi/pip integration is used
+instead of the legacy Python scripts.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+::::
+
 ::::{envvar} RULES_PYTHON_EXTRACT_ROOT
 
 Directory to use as the root for creating files necessary for bootstrapping so
diff --git a/python/private/internal_config_repo.bzl b/python/private/internal_config_repo.bzl
@@ -20,6 +20,8 @@ settings for rules to later use.
 
 load(":repo_utils.bzl", "repo_utils")
 
+_ENABLE_PIPSTAR_ENVVAR_NAME = "RULES_PYTHON_ENABLE_PIPSTAR"
+_ENABLE_PIPSTAR_DEFAULT = "0"
 _ENABLE_PYSTAR_ENVVAR_NAME = "RULES_PYTHON_ENABLE_PYSTAR"
 _ENABLE_PYSTAR_DEFAULT = "1"
 _ENABLE_DEPRECATION_WARNINGS_ENVVAR_NAME = "RULES_PYTHON_DEPRECATION_WARNINGS"
@@ -28,6 +30,7 @@ _ENABLE_DEPRECATION_WARNINGS_DEFAULT = "0"
 _CONFIG_TEMPLATE = """\
 config = struct(
   enable_pystar = {enable_pystar},
+  enable_pipstar = {enable_pipstar},
   enable_deprecation_warnings = {enable_deprecation_warnings},
   BuiltinPyInfo = getattr(getattr(native, "legacy_globals", None), "PyInfo", {builtin_py_info_symbol}),
   BuiltinPyRuntimeInfo = getattr(getattr(native, "legacy_globals", None), "PyRuntimeInfo", {builtin_py_runtime_info_symbol}),
@@ -84,6 +87,7 @@ def _internal_config_repo_impl(rctx):
 
     rctx.file("rules_python_config.bzl", _CONFIG_TEMPLATE.format(
         enable_pystar = enable_pystar,
+        enable_pipstar = _bool_from_environ(rctx, _ENABLE_PIPSTAR_ENVVAR_NAME, _ENABLE_PIPSTAR_DEFAULT),
         enable_deprecation_warnings = _bool_from_environ(rctx, _ENABLE_DEPRECATION_WARNINGS_ENVVAR_NAME, _ENABLE_DEPRECATION_WARNINGS_DEFAULT),
         builtin_py_info_symbol = builtin_py_info_symbol,
         builtin_py_runtime_info_symbol = builtin_py_runtime_info_symbol,
diff --git a/python/private/pypi/whl_installer/arguments.py b/python/private/pypi/whl_installer/arguments.py
@@ -47,6 +47,11 @@ def parser(**kwargs: Any) -> argparse.ArgumentParser:
         type=Platform.from_string,
         help="Platforms to target dependencies. Can be used multiple times.",
     )
+    parser.add_argument(
+        "--enable-pipstar",
+        action="store_true",
+        help="Disable certain code paths if we expect to process the whl in Starlark.",
+    )
     parser.add_argument(
         "--pip_data_exclude",
         action="store",
diff --git a/python/private/pypi/whl_installer/wheel_installer.py b/python/private/pypi/whl_installer/wheel_installer.py
@@ -104,6 +104,7 @@ def _setup_namespace_pkg_compatibility(wheel_dir: str) -> None:
 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("."),
@@ -114,6 +115,7 @@ def _extract_wheel(
         wheel_file: the filepath of the .whl
         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
     """
 
@@ -123,26 +125,31 @@ def _extract_wheel(
     if not enable_implicit_namespace_pkgs:
         _setup_namespace_pkg_compatibility(installation_dir)
 
-    extras_requested = extras[whl.name] if whl.name in extras else set()
-
-    dependencies = whl.dependencies(extras_requested, platforms)
+    metadata = {
+        "python_version": f"{sys.version_info[0]}.{sys.version_info[1]}.{sys.version_info[2]}",
+        "entry_points": [
+            {
+                "name": name,
+                "module": module,
+                "attribute": attribute,
+            }
+            for name, (module, attribute) in sorted(whl.entry_points().items())
+        ],
+    }
+    if not enable_pipstar:
+        extras_requested = extras[whl.name] if whl.name in extras else set()
+        dependencies = whl.dependencies(extras_requested, platforms)
+
+        metadata.update(
+            {
+                "name": whl.name,
+                "version": whl.version,
+                "deps": dependencies.deps,
+                "deps_by_platform": dependencies.deps_select,
+            }
+        )
 
     with open(os.path.join(installation_dir, "metadata.json"), "w") as f:
-        metadata = {
-            "name": whl.name,
-            "version": whl.version,
-            "deps": dependencies.deps,
-            "python_version": f"{sys.version_info[0]}.{sys.version_info[1]}.{sys.version_info[2]}",
-            "deps_by_platform": dependencies.deps_select,
-            "entry_points": [
-                {
-                    "name": name,
-                    "module": module,
-                    "attribute": attribute,
-                }
-                for name, (module, attribute) in sorted(whl.entry_points().items())
-            ],
-        }
         json.dump(metadata, f)
 
 
@@ -161,6 +168,7 @@ def main() -> None:
         _extract_wheel(
             wheel_file=whl,
             extras=extras,
+            enable_pipstar=args.enable_pipstar,
             enable_implicit_namespace_pkgs=args.enable_implicit_namespace_pkgs,
             platforms=arguments.get_platforms(args),
         )
diff --git a/python/private/pypi/whl_library.bzl b/python/private/pypi/whl_library.bzl
@@ -14,16 +14,19 @@
 
 ""
 
+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:repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")
 load(":attrs.bzl", "ATTRS", "use_isolated")
 load(":deps.bzl", "all_repo_names", "record_files")
 load(":generate_whl_library_build_bazel.bzl", "generate_whl_library_build_bazel")
+load(":parse_requirements.bzl", "host_platform")
 load(":parse_whl_name.bzl", "parse_whl_name")
 load(":patch_whl.bzl", "patch_whl")
 load(":pypi_repo_utils.bzl", "pypi_repo_utils")
+load(":whl_metadata.bzl", "whl_metadata")
 load(":whl_target_platforms.bzl", "whl_target_platforms")
 
 _CPPFLAGS = "CPPFLAGS"
@@ -340,79 +343,147 @@ def _whl_library_impl(rctx):
                 timeout = rctx.attr.timeout,
             )
 
-    target_platforms = rctx.attr.experimental_target_platforms or []
-    if target_platforms:
-        parsed_whl = parse_whl_name(whl_path.basename)
-
-        # NOTE @aignas 2023-12-04: if the wheel is a platform specific wheel, we
-        # only include deps for that target platform
-        if parsed_whl.platform_tag != "any":
-            target_platforms = [
-                p.target_platform
-                for p in whl_target_platforms(
-                    platform_tag = parsed_whl.platform_tag,
-                    abi_tag = parsed_whl.abi_tag.strip("tm"),
-                )
-            ]
-
-    pypi_repo_utils.execute_checked(
-        rctx,
-        op = "whl_library.ExtractWheel({}, {})".format(rctx.attr.name, whl_path),
-        python = python_interpreter,
-        arguments = args + [
-            "--whl-file",
-            whl_path,
-        ] + ["--platform={}".format(p) for p in target_platforms],
-        srcs = rctx.attr._python_srcs,
-        environment = environment,
-        quiet = rctx.attr.quiet,
-        timeout = rctx.attr.timeout,
-        logger = logger,
-    )
+    if rp_config.enable_pipstar:
+        pypi_repo_utils.execute_checked(
+            rctx,
+            op = "whl_library.ExtractWheel({}, {})".format(rctx.attr.name, whl_path),
+            python = python_interpreter,
+            arguments = args + [
+                "--whl-file",
+                whl_path,
+                "--enable-pipstar",
+            ],
+            srcs = rctx.attr._python_srcs,
+            environment = environment,
+            quiet = rctx.attr.quiet,
+            timeout = rctx.attr.timeout,
+            logger = logger,
+        )
 
-    metadata = json.decode(rctx.read("metadata.json"))
-    rctx.delete("metadata.json")
+        metadata = json.decode(rctx.read("metadata.json"))
+        rctx.delete("metadata.json")
+        python_version = metadata["python_version"]
 
-    # NOTE @aignas 2024-06-22: this has to live on until we stop supporting
-    # passing `twine` as a `:pkg` library via the `WORKSPACE` builds.
-    #
-    # See ../../packaging.bzl line 190
-    entry_points = {}
-    for item in metadata["entry_points"]:
-        name = item["name"]
-        module = item["module"]
-        attribute = item["attribute"]
-
-        # There is an extreme edge-case with entry_points that end with `.py`
-        # See: https://github.com/bazelbuild/bazel/blob/09c621e4cf5b968f4c6cdf905ab142d5961f9ddc/src/test/java/com/google/devtools/build/lib/rules/python/PyBinaryConfiguredTargetTest.java#L174
-        entry_point_without_py = name[:-3] + "_py" if name.endswith(".py") else name
-        entry_point_target_name = (
-            _WHEEL_ENTRY_POINT_PREFIX + "_" + entry_point_without_py
+        # NOTE @aignas 2024-06-22: this has to live on until we stop supporting
+        # passing `twine` as a `:pkg` library via the `WORKSPACE` builds.
+        #
+        # See ../../packaging.bzl line 190
+        entry_points = {}
+        for item in metadata["entry_points"]:
+            name = item["name"]
+            module = item["module"]
+            attribute = item["attribute"]
+
+            # There is an extreme edge-case with entry_points that end with `.py`
+            # See: https://github.com/bazelbuild/bazel/blob/09c621e4cf5b968f4c6cdf905ab142d5961f9ddc/src/test/java/com/google/devtools/build/lib/rules/python/PyBinaryConfiguredTargetTest.java#L174
+            entry_point_without_py = name[:-3] + "_py" if name.endswith(".py") else name
+            entry_point_target_name = (
+                _WHEEL_ENTRY_POINT_PREFIX + "_" + entry_point_without_py
+            )
+            entry_point_script_name = entry_point_target_name + ".py"
+
+            rctx.file(
+                entry_point_script_name,
+                _generate_entry_point_contents(module, attribute),
+            )
+            entry_points[entry_point_without_py] = entry_point_script_name
+
+        metadata = whl_metadata(
+            install_dir = whl_path.dirname.get_child("site-packages"),
+            read_fn = rctx.read,
+            logger = logger,
         )
-        entry_point_script_name = entry_point_target_name + ".py"
 
-        rctx.file(
-            entry_point_script_name,
-            _generate_entry_point_contents(module, attribute),
+        build_file_contents = generate_whl_library_build_bazel(
+            name = whl_path.basename,
+            dep_template = rctx.attr.dep_template or "@{}{{name}}//:{{target}}".format(rctx.attr.repo_prefix),
+            entry_points = entry_points,
+            metadata_name = metadata.name,
+            metadata_version = metadata.version,
+            default_python_version = python_version,
+            requires_dist = metadata.requires_dist,
+            target_platforms = rctx.attr.experimental_target_platforms or [host_platform(rctx)],
+            # 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,
+            group_deps = rctx.attr.group_deps,
+            group_name = rctx.attr.group_name,
         )
-        entry_points[entry_point_without_py] = entry_point_script_name
-
-    build_file_contents = generate_whl_library_build_bazel(
-        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-04-14: load through the hub:
-        dependencies = metadata["deps"],
-        dependencies_by_platform = metadata["deps_by_platform"],
-        annotation = None if not rctx.attr.annotation else struct(**json.decode(rctx.read(rctx.attr.annotation))),
-        data_exclude = rctx.attr.pip_data_exclude,
-        group_deps = rctx.attr.group_deps,
-        group_name = rctx.attr.group_name,
-        tags = [
-            "pypi_name={}".format(metadata["name"]),
-            "pypi_version={}".format(metadata["version"]),
-        ],
-    )
+    else:
+        target_platforms = rctx.attr.experimental_target_platforms or []
+        if target_platforms:
+            parsed_whl = parse_whl_name(whl_path.basename)
+
+            # NOTE @aignas 2023-12-04: if the wheel is a platform specific wheel, we
+            # only include deps for that target platform
+            if parsed_whl.platform_tag != "any":
+                target_platforms = [
+                    p.target_platform
+                    for p in whl_target_platforms(
+                        platform_tag = parsed_whl.platform_tag,
+                        abi_tag = parsed_whl.abi_tag.strip("tm"),
+                    )
+                ]
+
+        pypi_repo_utils.execute_checked(
+            rctx,
+            op = "whl_library.ExtractWheel({}, {})".format(rctx.attr.name, whl_path),
+            python = python_interpreter,
+            arguments = args + [
+                "--whl-file",
+                whl_path,
+            ] + ["--platform={}".format(p) for p in target_platforms],
+            srcs = rctx.attr._python_srcs,
+            environment = environment,
+            quiet = rctx.attr.quiet,
+            timeout = rctx.attr.timeout,
+            logger = logger,
+        )
+
+        metadata = json.decode(rctx.read("metadata.json"))
+        rctx.delete("metadata.json")
+
+        # NOTE @aignas 2024-06-22: this has to live on until we stop supporting
+        # passing `twine` as a `:pkg` library via the `WORKSPACE` builds.
+        #
+        # See ../../packaging.bzl line 190
+        entry_points = {}
+        for item in metadata["entry_points"]:
+            name = item["name"]
+            module = item["module"]
+            attribute = item["attribute"]
+
+            # There is an extreme edge-case with entry_points that end with `.py`
+            # See: https://github.com/bazelbuild/bazel/blob/09c621e4cf5b968f4c6cdf905ab142d5961f9ddc/src/test/java/com/google/devtools/build/lib/rules/python/PyBinaryConfiguredTargetTest.java#L174
+            entry_point_without_py = name[:-3] + "_py" if name.endswith(".py") else name
+            entry_point_target_name = (
+                _WHEEL_ENTRY_POINT_PREFIX + "_" + entry_point_without_py
+            )
+            entry_point_script_name = entry_point_target_name + ".py"
+
+            rctx.file(
+                entry_point_script_name,
+                _generate_entry_point_contents(module, attribute),
+            )
+            entry_points[entry_point_without_py] = entry_point_script_name
+
+        build_file_contents = generate_whl_library_build_bazel(
+            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-04-14: load through the hub:
+            dependencies = metadata["deps"],
+            dependencies_by_platform = metadata["deps_by_platform"],
+            annotation = None if not rctx.attr.annotation else struct(**json.decode(rctx.read(rctx.attr.annotation))),
+            data_exclude = rctx.attr.pip_data_exclude,
+            group_deps = rctx.attr.group_deps,
+            group_name = rctx.attr.group_name,
+            tags = [
+                "pypi_name={}".format(metadata["name"]),
+                "pypi_version={}".format(metadata["version"]),
+            ],
+        )
+
     rctx.file("BUILD.bazel", build_file_contents)
 
     return
diff --git a/tests/pypi/whl_installer/wheel_installer_test.py b/tests/pypi/whl_installer/wheel_installer_test.py
@@ -72,6 +72,7 @@ def test_wheel_exists(self) -> None:
             extras={},
             enable_implicit_namespace_pkgs=False,
             platforms=[],
+            enable_pipstar = False,
         )
 
         want_files = [

Original file line number	Diff line number	Diff line change
`@@ -72,6 +72,7 @@ def test_wheel_exists(self) -> None:`
`72`	`72`	`extras={},`
`73`	`73`	`enable_implicit_namespace_pkgs=False,`
`74`	`74`	`platforms=[],`
	`75`	`+ enable_pipstar = False,`
`75`	`76`	`)`
`76`	`77`
`77`	`78`	`want_files = [`