From a5dc20c558948db900dd0964ac23dcd2787d946c Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Sun, 6 Jul 2025 22:13:46 +0900
Subject: [PATCH 1/4] refactor(pypi): move the platform config to MODULE.bazel

This splits the configuration that we have for the `rules_python` and
the defaults that we set for our users from the actual unit tests where
we check that the extension is working correctly.

With this we will be able to dog food the API and point users into the
`MODULE.bazel` as the example snippet.

Work towards #2949
Work towards #2747
---
 MODULE.bazel                             |  56 ++++++++++
 python/private/pypi/extension.bzl        |  60 +----------
 tests/pypi/extension/extension_tests.bzl | 132 ++++++++++-------------
 3 files changed, 115 insertions(+), 133 deletions(-)

diff --git a/MODULE.bazel b/MODULE.bazel
index b1d8711815..a9b51951b2 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -60,6 +60,62 @@ register_toolchains("@pythons_hub//:all")
 # Install twine for our own runfiles wheel publishing and allow bzlmod users to use it.
 
 pip = use_extension("//python/extensions:pip.bzl", "pip")
+
+# NOTE @aignas 2025-07-06: we define these platforms to keep backwards compatibility with the
+# current `experimental_index_url` implementation. Whilst we stabilize the API this list may be
+# updated with a mention in the CHANGELOG.
+[
+    pip.default(
+        arch_name = cpu,
+        config_settings = [
+            "@platforms//cpu:{}".format(cpu),
+            "@platforms//os:linux",
+        ],
+        env = {"platform_version": "0"},
+        os_name = "linux",
+        platform = "linux_{}".format(cpu),
+    )
+    for cpu in [
+        "x86_64",
+        "aarch64",
+        # TODO @aignas 2025-05-19: only leave tier 0-1 cpus when stabilizing the
+        # `pip.default` extension. i.e. drop the below values - users will have to
+        # define themselves if they need them.
+        "arm",
+        "ppc",
+        "s390x",
+    ]
+]
+
+[
+    pip.default(
+        arch_name = cpu,
+        config_settings = [
+            "@platforms//cpu:{}".format(cpu),
+            "@platforms//os:osx",
+        ],
+        # We choose the oldest non-EOL version at the time when we release `rules_python`.
+        # See https://endoflife.date/macos
+        env = {"platform_version": "14.0"},
+        os_name = "osx",
+        platform = "osx_{}".format(cpu),
+    )
+    for cpu in [
+        "aarch64",
+        "x86_64",
+    ]
+]
+
+pip.default(
+    arch_name = "x86_64",
+    config_settings = [
+        "@platforms//cpu:x86_64",
+        "@platforms//os:windows",
+    ],
+    env = {"platform_version": "0"},
+    os_name = "windows",
+    platform = "windows_x86_64",
+)
 pip.parse(
     # NOTE @aignas 2024-10-26: We have an integration test that depends on us
     # being able to build sdists for this hub, so explicitly set this to False.
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index a0095f8f15..505458008f 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -393,64 +393,6 @@ def _configure(config, *, platform, os_name, arch_name, config_settings, env = {
     else:
         config["platforms"].pop(platform)
 
-def _create_config(defaults):
-    if defaults["platforms"]:
-        return struct(**defaults)
-
-    # NOTE: We have this so that it is easier to maintain unit tests assuming certain
-    # defaults
-    for cpu in [
-        "x86_64",
-        "aarch64",
-        # TODO @aignas 2025-05-19: only leave tier 0-1 cpus when stabilizing the
-        # `pip.default` extension. i.e. drop the below values - users will have to
-        # define themselves if they need them.
-        "arm",
-        "ppc",
-        "s390x",
-    ]:
-        _configure(
-            defaults,
-            arch_name = cpu,
-            os_name = "linux",
-            platform = "linux_{}".format(cpu),
-            config_settings = [
-                "@platforms//os:linux",
-                "@platforms//cpu:{}".format(cpu),
-            ],
-            env = {"platform_version": "0"},
-        )
-    for cpu in [
-        "aarch64",
-        "x86_64",
-    ]:
-        _configure(
-            defaults,
-            arch_name = cpu,
-            # We choose the oldest non-EOL version at the time when we release `rules_python`.
-            # See https://endoflife.date/macos
-            os_name = "osx",
-            platform = "osx_{}".format(cpu),
-            config_settings = [
-                "@platforms//os:osx",
-                "@platforms//cpu:{}".format(cpu),
-            ],
-            env = {"platform_version": "14.0"},
-        )
-
-    _configure(
-        defaults,
-        arch_name = "x86_64",
-        os_name = "windows",
-        platform = "windows_x86_64",
-        config_settings = [
-            "@platforms//os:windows",
-            "@platforms//cpu:x86_64",
-        ],
-        env = {"platform_version": "0"},
-    )
-    return struct(**defaults)
-
 def parse_modules(
         module_ctx,
         _fail = fail,
@@ -527,7 +469,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
                 # for what. We could also model the `cp313t` freethreaded as separate platforms.
             )
 
-    config = _create_config(defaults)
+    config = struct(**defaults)
 
     # TODO @aignas 2025-06-03: Merge override API with the builder?
     _overriden_whl_set = {}
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index 146293ee8d..cf96d4005a 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -56,7 +56,23 @@ def _mod(*, name, default = [], parse = [], override = [], whl_mods = [], is_roo
             parse = parse,
             override = override,
             whl_mods = whl_mods,
-            default = default,
+            default = default or [
+                _default(
+                    platform = "{}_{}".format(os, cpu),
+                    os_name = os,
+                    arch_name = cpu,
+                    config_settings = [
+                        "@platforms//os:{}".format(os),
+                        "@platforms//cpu:{}".format(cpu),
+                    ],
+                )
+                for os, cpu in [
+                    ("linux", "x86_64"),
+                    ("linux", "aarch64"),
+                    ("osx", "aarch64"),
+                    ("windows", "aarch64"),
+                ]
+            ],
         ),
         is_root = is_root,
     )
@@ -235,19 +251,18 @@ def _test_simple_multiple_requirements(env):
     pypi.hub_group_map().contains_exactly({"pypi": {}})
     pypi.hub_whl_map().contains_exactly({"pypi": {
         "simple": {
-            "pypi_315_simple_osx_aarch64_osx_x86_64": [
+            "pypi_315_simple_osx_aarch64": [
                 whl_config_setting(
                     target_platforms = [
                         "cp315_osx_aarch64",
-                        "cp315_osx_x86_64",
                     ],
                     version = "3.15",
                 ),
             ],
-            "pypi_315_simple_windows_x86_64": [
+            "pypi_315_simple_windows_aarch64": [
                 whl_config_setting(
                     target_platforms = [
-                        "cp315_windows_x86_64",
+                        "cp315_windows_aarch64",
                     ],
                     version = "3.15",
                 ),
@@ -255,12 +270,12 @@ def _test_simple_multiple_requirements(env):
         },
     }})
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_simple_osx_aarch64_osx_x86_64": {
+        "pypi_315_simple_osx_aarch64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "simple==0.0.2 --hash=sha256:deadb00f",
         },
-        "pypi_315_simple_windows_x86_64": {
+        "pypi_315_simple_windows_aarch64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "simple==0.0.1 --hash=sha256:deadbeef",
@@ -310,24 +325,20 @@ torch==2.4.1 ; platform_machine != 'x86_64' \
     pypi.hub_group_map().contains_exactly({"pypi": {}})
     pypi.hub_whl_map().contains_exactly({"pypi": {
         "torch": {
-            "pypi_315_torch_linux_aarch64_linux_arm_linux_ppc_linux_s390x_osx_aarch64": [
+            "pypi_315_torch_linux_aarch64_osx_aarch64_windows_aarch64": [
                 whl_config_setting(
                     target_platforms = [
                         "cp315_linux_aarch64",
-                        "cp315_linux_arm",
-                        "cp315_linux_ppc",
-                        "cp315_linux_s390x",
                         "cp315_osx_aarch64",
+                        "cp315_windows_aarch64",
                     ],
                     version = "3.15",
                 ),
             ],
-            "pypi_315_torch_linux_x86_64_osx_x86_64_windows_x86_64": [
+            "pypi_315_torch_linux_x86_64": [
                 whl_config_setting(
                     target_platforms = [
                         "cp315_linux_x86_64",
-                        "cp315_osx_x86_64",
-                        "cp315_windows_x86_64",
                     ],
                     version = "3.15",
                 ),
@@ -335,12 +346,12 @@ torch==2.4.1 ; platform_machine != 'x86_64' \
         },
     }})
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_torch_linux_aarch64_linux_arm_linux_ppc_linux_s390x_osx_aarch64": {
+        "pypi_315_torch_linux_aarch64_osx_aarch64_windows_aarch64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1 --hash=sha256:deadbeef",
         },
-        "pypi_315_torch_linux_x86_64_osx_x86_64_windows_x86_64": {
+        "pypi_315_torch_linux_x86_64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1+cpu",
@@ -385,6 +396,23 @@ def _test_torch_experimental_index_url(env):
         module_ctx = _mock_mctx(
             _mod(
                 name = "rules_python",
+                default = [
+                    _default(
+                        platform = "{}_{}".format(os, cpu),
+                        os_name = os,
+                        arch_name = cpu,
+                        config_settings = [
+                            "@platforms//os:{}".format(os),
+                            "@platforms//cpu:{}".format(cpu),
+                        ],
+                    )
+                    for os, cpu in [
+                        ("linux", "aarch64"),
+                        ("linux", "x86_64"),
+                        ("osx", "aarch64"),
+                        ("windows", "x86_64"),
+                    ]
+                ],
                 parse = [
                     _parse(
                         hub_name = "pypi",
@@ -444,34 +472,26 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
     pypi.hub_whl_map().contains_exactly({"pypi": {
         "torch": {
             "pypi_312_torch_cp312_cp312_linux_x86_64_8800deef": [
-                struct(
-                    config_setting = None,
+                whl_config_setting(
                     filename = "torch-2.4.1+cpu-cp312-cp312-linux_x86_64.whl",
-                    target_platforms = None,
                     version = "3.12",
                 ),
             ],
             "pypi_312_torch_cp312_cp312_manylinux_2_17_aarch64_36109432": [
-                struct(
-                    config_setting = None,
+                whl_config_setting(
                     filename = "torch-2.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
-                    target_platforms = None,
                     version = "3.12",
                 ),
             ],
             "pypi_312_torch_cp312_cp312_win_amd64_3a570e5c": [
-                struct(
-                    config_setting = None,
+                whl_config_setting(
                     filename = "torch-2.4.1+cpu-cp312-cp312-win_amd64.whl",
-                    target_platforms = None,
                     version = "3.12",
                 ),
             ],
             "pypi_312_torch_cp312_none_macosx_11_0_arm64_72b484d5": [
-                struct(
-                    config_setting = None,
+                whl_config_setting(
                     filename = "torch-2.4.1-cp312-none-macosx_11_0_arm64.whl",
-                    target_platforms = None,
                     version = "3.12",
                 ),
             ],
@@ -482,7 +502,6 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_x86_64",
-                "osx_x86_64",
                 "windows_x86_64",
             ],
             "filename": "torch-2.4.1+cpu-cp312-cp312-linux_x86_64.whl",
@@ -495,9 +514,6 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "osx_aarch64",
             ],
             "filename": "torch-2.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
@@ -510,7 +526,6 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_x86_64",
-                "osx_x86_64",
                 "windows_x86_64",
             ],
             "filename": "torch-2.4.1+cpu-cp312-cp312-win_amd64.whl",
@@ -523,9 +538,6 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "osx_aarch64",
             ],
             "filename": "torch-2.4.1-cp312-none-macosx_11_0_arm64.whl",
@@ -817,13 +829,9 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "linux_x86_64",
                 "osx_aarch64",
-                "osx_x86_64",
-                "windows_x86_64",
+                "windows_aarch64",
             ],
             "extra_pip_args": ["--extra-args-for-sdist-building"],
             "filename": "any-name.tar.gz",
@@ -836,13 +844,9 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "linux_x86_64",
                 "osx_aarch64",
-                "osx_x86_64",
-                "windows_x86_64",
+                "windows_aarch64",
             ],
             "filename": "direct_without_sha-0.0.1-py3-none-any.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
@@ -866,13 +870,9 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "linux_x86_64",
                 "osx_aarch64",
-                "osx_x86_64",
-                "windows_x86_64",
+                "windows_aarch64",
             ],
             "filename": "simple-0.0.1-py3-none-any.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
@@ -884,13 +884,9 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "linux_x86_64",
                 "osx_aarch64",
-                "osx_x86_64",
-                "windows_x86_64",
+                "windows_aarch64",
             ],
             "extra_pip_args": ["--extra-args-for-sdist-building"],
             "filename": "simple-0.0.1.tar.gz",
@@ -903,13 +899,9 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "linux_x86_64",
                 "osx_aarch64",
-                "osx_x86_64",
-                "windows_x86_64",
+                "windows_aarch64",
             ],
             "filename": "some_pkg-0.0.1-py3-none-any.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
@@ -921,13 +913,9 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
                 "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
                 "linux_x86_64",
                 "osx_aarch64",
-                "osx_x86_64",
-                "windows_x86_64",
+                "windows_aarch64",
             ],
             "filename": "some-other-pkg-0.0.1-py3-none-any.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
@@ -995,26 +983,22 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     pypi.hub_whl_map().contains_exactly({
         "pypi": {
             "optimum": {
-                "pypi_315_optimum_linux_aarch64_linux_arm_linux_ppc_linux_s390x_linux_x86_64": [
+                "pypi_315_optimum_linux_aarch64_linux_x86_64": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
                             "cp315_linux_aarch64",
-                            "cp315_linux_arm",
-                            "cp315_linux_ppc",
-                            "cp315_linux_s390x",
                             "cp315_linux_x86_64",
                         ],
                         config_setting = None,
                         filename = None,
                     ),
                 ],
-                "pypi_315_optimum_osx_aarch64_osx_x86_64": [
+                "pypi_315_optimum_osx_aarch64": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
                             "cp315_osx_aarch64",
-                            "cp315_osx_x86_64",
                         ],
                         config_setting = None,
                         filename = None,
@@ -1025,12 +1009,12 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     })
 
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_optimum_linux_aarch64_linux_arm_linux_ppc_linux_s390x_linux_x86_64": {
+        "pypi_315_optimum_linux_aarch64_linux_x86_64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime-gpu]==1.17.1",
         },
-        "pypi_315_optimum_osx_aarch64_osx_x86_64": {
+        "pypi_315_optimum_osx_aarch64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime]==1.17.1",

From ead607103cf40fbe6bd1d26fa89c976fd0f78c99 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Fri, 4 Jul 2025 09:59:48 +0900
Subject: [PATCH 2/4] fix(pypi): correctly handle custom names in pipstar
 platforms

Before it seems that we were relying on particular names in the pipstar
platforms. This ensures that we rely on this less. Whilst at it fix a
few typos and improve the formatting of the code.

Work towards #2949
Work towards #2747
---
 CHANGELOG.md                               |  2 ++
 python/private/pypi/BUILD.bazel            |  4 ----
 python/private/pypi/evaluate_markers.bzl   |  2 +-
 python/private/pypi/extension.bzl          | 14 ++++++++++++--
 python/private/pypi/parse_requirements.bzl |  6 +++++-
 python/private/pypi/pep508_env.bzl         |  5 -----
 tests/pypi/extension/extension_tests.bzl   | 20 +++++++++-----------
 tests/pypi/pep508/evaluate_tests.bzl       |  3 ++-
 8 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 81768af36a..2b57af606e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -82,6 +82,8 @@ END_UNRELEASED_TEMPLATE
 * (toolchains) `local_runtime_repo` now checks if the include directory exists
   before attempting to watch it, fixing issues on macOS with system Python
   ({gh-issue}`3043`).
+* (pypi) The pipstar `defaults` configuration now supports any custom platform
+  name.
 
 {#v0-0-0-added}
 ### Added
diff --git a/python/private/pypi/BUILD.bazel b/python/private/pypi/BUILD.bazel
index 2666197786..b098f29e94 100644
--- a/python/private/pypi/BUILD.bazel
+++ b/python/private/pypi/BUILD.bazel
@@ -252,10 +252,6 @@ bzl_library(
 bzl_library(
     name = "pep508_env_bzl",
     srcs = ["pep508_env.bzl"],
-    deps = [
-        ":pep508_platform_bzl",
-        "//python/private:version_bzl",
-    ],
 )
 
 bzl_library(
diff --git a/python/private/pypi/evaluate_markers.bzl b/python/private/pypi/evaluate_markers.bzl
index 2b805c33e6..6167cdbc96 100644
--- a/python/private/pypi/evaluate_markers.bzl
+++ b/python/private/pypi/evaluate_markers.bzl
@@ -57,7 +57,7 @@ def evaluate_markers_py(mrctx, *, requirements, python_interpreter, python_inter
 
     Args:
         mrctx: repository_ctx or module_ctx.
-        requirements: list[str] of the requirement file lines to evaluate.
+        requirements: {type}`dict[str, list[str]]` of the requirement file lines to evaluate.
         python_interpreter: str, path to the python_interpreter to use to
             evaluate the env markers in the given requirements files. It will
             be only called if the requirements files have env markers. This
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index a0095f8f15..64006bd9db 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -76,7 +76,11 @@ def _platforms(*, python_version, minor_mapping, config):
 
     for platform, values in config.platforms.items():
         key = "{}_{}".format(abi, platform)
-        platforms[key] = env(key) | values.env
+        platforms[key] = env(struct(
+            abi = abi,
+            os = values.os_name,
+            arch = values.arch_name,
+        )) | values.env
     return platforms
 
 def _create_whl_repos(
@@ -348,7 +352,7 @@ def _whl_repo(*, src, whl_library_args, is_multiple_versions, download_only, net
     args["filename"] = src.filename
     if not enable_pipstar:
         args["experimental_target_platforms"] = [
-            # Get rid of the version fot the target platforms because we are
+            # Get rid of the version for the target platforms because we are
             # passing the interpreter any way. Ideally we should search of ways
             # how to pass the target platforms through the hub repo.
             p.partition("_")[2]
@@ -383,6 +387,12 @@ def _configure(config, *, platform, os_name, arch_name, config_settings, env = {
             if key not in _SUPPORTED_PEP508_KEYS:
                 fail("Unsupported key in the PEP508 environment: {}".format(key))
 
+        if not os_name:
+            fail("'os_name' is required")
+
+        if not arch_name:
+            fail("'arch_name' is required")
+
         config["platforms"][platform] = struct(
             name = platform.replace("-", "_").lower(),
             os_name = os_name,
diff --git a/python/private/pypi/parse_requirements.bzl b/python/private/pypi/parse_requirements.bzl
index e4a8b90acb..9c610f11d3 100644
--- a/python/private/pypi/parse_requirements.bzl
+++ b/python/private/pypi/parse_requirements.bzl
@@ -402,6 +402,10 @@ def _add_dists(*, requirement, index_urls, logger = None):
         ]))
 
     # Filter out the wheels that are incompatible with the target_platforms.
-    whls = select_whls(whls = whls, want_platforms = requirement.target_platforms, logger = logger)
+    whls = select_whls(
+        whls = whls,
+        want_platforms = requirement.target_platforms,
+        logger = logger,
+    )
 
     return whls, sdist
diff --git a/python/private/pypi/pep508_env.bzl b/python/private/pypi/pep508_env.bzl
index a6efb3c50c..c2d404bc3e 100644
--- a/python/private/pypi/pep508_env.bzl
+++ b/python/private/pypi/pep508_env.bzl
@@ -15,8 +15,6 @@
 """This module is for implementing PEP508 environment definition.
 """
 
-load(":pep508_platform.bzl", "platform_from_str")
-
 # See https://stackoverflow.com/a/45125525
 platform_machine_aliases = {
     # These pairs mean the same hardware, but different values may be used
@@ -175,9 +173,6 @@ def env(target_platform, *, extra = None):
     if extra != None:
         env["extra"] = extra
 
-    if type(target_platform) == type(""):
-        target_platform = platform_from_str(target_platform, python_version = "")
-
     if target_platform.abi:
         minor_version, _, micro_version = target_platform.abi[3:].partition(".")
         micro_version = micro_version or "0"
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index 146293ee8d..51e11342fc 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -1048,7 +1048,9 @@ def _test_pipstar_platforms(env):
                 name = "rules_python",
                 default = [
                     _default(
-                        platform = "{}_{}".format(os, cpu),
+                        platform = "my{}_{}".format(os, cpu),
+                        os_name = os,
+                        arch_name = cpu,
                         config_settings = [
                             "@platforms//os:{}".format(os),
                             "@platforms//cpu:{}".format(cpu),
@@ -1086,24 +1088,20 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     pypi.hub_whl_map().contains_exactly({
         "pypi": {
             "optimum": {
-                "pypi_315_optimum_linux_x86_64": [
+                "pypi_315_optimum_mylinux_x86_64": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
-                            "cp315_linux_x86_64",
+                            "cp315_mylinux_x86_64",
                         ],
-                        config_setting = None,
-                        filename = None,
                     ),
                 ],
-                "pypi_315_optimum_osx_aarch64": [
+                "pypi_315_optimum_myosx_aarch64": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
-                            "cp315_osx_aarch64",
+                            "cp315_myosx_aarch64",
                         ],
-                        config_setting = None,
-                        filename = None,
                     ),
                 ],
             },
@@ -1111,12 +1109,12 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     })
 
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_optimum_linux_x86_64": {
+        "pypi_315_optimum_mylinux_x86_64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime-gpu]==1.17.1",
         },
-        "pypi_315_optimum_osx_aarch64": {
+        "pypi_315_optimum_myosx_aarch64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime]==1.17.1",
diff --git a/tests/pypi/pep508/evaluate_tests.bzl b/tests/pypi/pep508/evaluate_tests.bzl
index 7b6c064b94..cc867f346c 100644
--- a/tests/pypi/pep508/evaluate_tests.bzl
+++ b/tests/pypi/pep508/evaluate_tests.bzl
@@ -16,6 +16,7 @@
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("//python/private/pypi:pep508_env.bzl", pep508_env = "env")  # buildifier: disable=bzl-visibility
 load("//python/private/pypi:pep508_evaluate.bzl", "evaluate", "tokenize")  # buildifier: disable=bzl-visibility
+load("//python/private/pypi:pep508_platform.bzl", "platform_from_str")  # buildifier: disable=bzl-visibility
 
 _tests = []
 
@@ -262,7 +263,7 @@ def _evaluate_with_aliases(env):
         },
     }.items():  # buildifier: @unsorted-dict-items
         for input, want in tests.items():
-            _check_evaluate(env, input, want, pep508_env(target_platform))
+            _check_evaluate(env, input, want, pep508_env(platform_from_str(target_platform, "")))
 
 _tests.append(_evaluate_with_aliases)
 

From 5c29535ff669ac965fc0aeb5578b254ee2c9c2a1 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 16 Jun 2025 22:52:43 +0900
Subject: [PATCH 3/4] fix(pypi): pull fewer wheels

Before this we would pull all of the wheels that the user target
configuration would be compatible with and that meant that it was not
customizable. This also meant that there were a lot of footguns in the
configuration where the select statements were not really foolproof.

With this PR we select only those sources that need to be for the
declared configurations.

Freethreaded support should be done by defining extre freethreaded
platforms using the new builder API.

Work towards #2747
Work towards #2759
Work towards #2849
---
 MODULE.bazel                                  |   5 +-
 examples/bzlmod/MODULE.bazel                  |  17 ++
 python/private/pypi/BUILD.bazel               |  13 +-
 python/private/pypi/evaluate_markers.bzl      |   8 +-
 python/private/pypi/extension.bzl             | 189 +++++++++----
 python/private/pypi/parse_requirements.bzl    |  94 ++++---
 python/private/pypi/select_whl.bzl            | 130 +++++++++
 python/private/pypi/whl_target_platforms.bzl  | 132 ---------
 tests/pypi/extension/extension_tests.bzl      | 255 +++++++++--------
 .../parse_requirements_tests.bzl              |  26 +-
 tests/pypi/select_whl/BUILD.bazel             |   3 +
 .../select_whl_tests.bzl                      | 262 ++++++++----------
 tests/pypi/whl_target_platforms/BUILD.bazel   |   3 -
 13 files changed, 620 insertions(+), 517 deletions(-)
 create mode 100644 python/private/pypi/select_whl.bzl
 create mode 100644 tests/pypi/select_whl/BUILD.bazel
 rename tests/pypi/{whl_target_platforms => select_whl}/select_whl_tests.bzl (57%)

diff --git a/MODULE.bazel b/MODULE.bazel
index b1d8711815..8643fceb7d 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -63,6 +63,8 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
 pip.parse(
     # NOTE @aignas 2024-10-26: We have an integration test that depends on us
     # being able to build sdists for this hub, so explicitly set this to False.
+    #
+    # how do we test sdists? Maybe just worth adding a single sdist somewhere?
     download_only = False,
     experimental_index_url = "https://pypi.org/simple",
     hub_name = "rules_python_publish_deps",
@@ -155,7 +157,6 @@ dev_pip = use_extension(
     dev_dependency = True,
 )
 dev_pip.parse(
-    download_only = True,
     experimental_index_url = "https://pypi.org/simple",
     hub_name = "dev_pip",
     parallel_download = False,
@@ -163,14 +164,12 @@ dev_pip.parse(
     requirements_lock = "//docs:requirements.txt",
 )
 dev_pip.parse(
-    download_only = True,
     experimental_index_url = "https://pypi.org/simple",
     hub_name = "dev_pip",
     python_version = "3.13",
     requirements_lock = "//docs:requirements.txt",
 )
 dev_pip.parse(
-    download_only = True,
     experimental_index_url = "https://pypi.org/simple",
     hub_name = "pypiserver",
     python_version = "3.11",
diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
index 841c096dcf..6b1ee2c351 100644
--- a/examples/bzlmod/MODULE.bazel
+++ b/examples/bzlmod/MODULE.bazel
@@ -158,6 +158,23 @@ pip.whl_mods(
 )
 use_repo(pip, "whl_mods_hub")
 
+# Because below we are using `windows_aarch64` platform, we have to define various
+# properties for it.
+pip.default(
+    arch_name = "aarch64",
+    config_settings = [
+        "@platforms//os:windows",
+        "@platforms//cpu:aarch64",
+    ],
+    env = {
+        "platform_version": "0",
+    },
+    os_name = "windows",
+    platform = "windows_aarch64",
+    platform_tags = ["win_amd64"],
+    want_abis = [],  # default to all ABIs
+)
+
 # To fetch pip dependencies, use pip.parse. We can pass in various options,
 # but typically we pass requirements and the Python version. The Python
 # version must have been configured by a corresponding `python.toolchain()`
diff --git a/python/private/pypi/BUILD.bazel b/python/private/pypi/BUILD.bazel
index b098f29e94..982c64436e 100644
--- a/python/private/pypi/BUILD.bazel
+++ b/python/private/pypi/BUILD.bazel
@@ -120,7 +120,6 @@ bzl_library(
         ":whl_config_setting_bzl",
         ":whl_library_bzl",
         ":whl_repo_name_bzl",
-        ":whl_target_platforms_bzl",
         "//python/private:full_version_bzl",
         "//python/private:normalize_name_bzl",
         "//python/private:version_bzl",
@@ -209,7 +208,7 @@ bzl_library(
         ":parse_requirements_txt_bzl",
         ":pypi_repo_utils_bzl",
         ":requirements_files_by_platform_bzl",
-        ":whl_target_platforms_bzl",
+        ":select_whl_bzl",
         "//python/private:normalize_name_bzl",
         "//python/private:repo_utils_bzl",
     ],
@@ -359,6 +358,15 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "select_whl_bzl",
+    srcs = ["select_whl.bzl"],
+    deps = [
+        ":parse_whl_name_bzl",
+        "//python/private:version_bzl",
+    ],
+)
+
 bzl_library(
     name = "simpleapi_download_bzl",
     srcs = ["simpleapi_download.bzl"],
@@ -422,5 +430,4 @@ bzl_library(
 bzl_library(
     name = "whl_target_platforms_bzl",
     srcs = ["whl_target_platforms.bzl"],
-    deps = [":parse_whl_name_bzl"],
 )
diff --git a/python/private/pypi/evaluate_markers.bzl b/python/private/pypi/evaluate_markers.bzl
index 6167cdbc96..ee8184ac3b 100644
--- a/python/private/pypi/evaluate_markers.bzl
+++ b/python/private/pypi/evaluate_markers.bzl
@@ -43,11 +43,11 @@ def evaluate_markers(*, requirements, platforms):
     for req_string, platform_strings in requirements.items():
         req = requirement(req_string)
         for platform_str in platform_strings:
-            env = platforms.get(platform_str)
-            if not env:
-                fail("Please define platform: '{}'".format(platform_str))
+            plat = platforms.get(platform_str)
+            if not plat:
+                continue
 
-            if evaluate(req.marker, env = env):
+            if evaluate(req.marker, env = plat.env):
                 ret.setdefault(req_string, []).append(platform_str)
 
     return ret
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index 64006bd9db..a5c1348060 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -72,15 +72,25 @@ def _platforms(*, python_version, minor_mapping, config):
         version = python_version,
         minor_mapping = minor_mapping,
     )
-    abi = "cp3{}".format(python_version[2:])
 
     for platform, values in config.platforms.items():
+        implementation = values.env["implementation_name"][:2].lower()
+        abi = "{}3{}".format(implementation, python_version[2:])
         key = "{}_{}".format(abi, platform)
-        platforms[key] = env(struct(
+
+        env_ = env(struct(
             abi = abi,
             os = values.os_name,
             arch = values.arch_name,
         )) | values.env
+        platforms[key] = struct(
+            env = env_,
+            want_abis = [
+                v.format(*python_version.split("."))
+                for v in values.want_abis
+            ],
+            platform_tags = values.platform_tags,
+        )
     return platforms
 
 def _create_whl_repos(
@@ -152,6 +162,8 @@ def _create_whl_repos(
             ))
         python_interpreter_target = available_interpreters[python_name]
 
+    # TODO @aignas 2025-06-29: we should not need the version in the pip_name if
+    # we are using pipstar and we are downloading the wheel using the downloader
     pip_name = "{}_{}".format(
         hub_name,
         version_label(pip_attr.python_version),
@@ -230,6 +242,11 @@ def _create_whl_repos(
             ),
             logger = logger,
         ),
+        platforms = _platforms(
+            python_version = pip_attr.python_version,
+            minor_mapping = minor_mapping,
+            config = config,
+        ),
         extra_pip_args = pip_attr.extra_pip_args,
         get_index_urls = get_index_urls,
         evaluate_markers = evaluate_markers,
@@ -359,27 +376,19 @@ def _whl_repo(*, src, whl_library_args, is_multiple_versions, download_only, net
             for p in src.target_platforms
         ]
 
-    # Pure python wheels or sdists may need to have a platform here
-    target_platforms = None
-    if is_whl and not src.filename.endswith("-any.whl"):
-        pass
-    elif is_multiple_versions:
-        target_platforms = src.target_platforms
-
     return struct(
         repo_name = whl_repo_name(src.filename, src.sha256),
         args = args,
         config_setting = whl_config_setting(
             version = python_version,
-            filename = src.filename,
-            target_platforms = target_platforms,
+            target_platforms = src.target_platforms,
         ),
     )
 
-def _configure(config, *, platform, os_name, arch_name, config_settings, env = {}, override = False):
+def _configure(config, *, platform, os_name, arch_name, config_settings, env = {}, want_abis, platform_tags, override = False):
     """Set the value in the config if the value is provided"""
     config.setdefault("platforms", {})
-    if platform:
+    if platform and (os_name or arch_name or config_settings or platform_tags or env):
         if not override and config.get("platforms", {}).get(platform):
             return
 
@@ -393,19 +402,35 @@ def _configure(config, *, platform, os_name, arch_name, config_settings, env = {
         if not arch_name:
             fail("'arch_name' is required")
 
+        if platform_tags and "any" not in platform_tags:
+            # the lowest priority one needs to be the first one
+            platform_tags = ["any"] + platform_tags
+
         config["platforms"][platform] = struct(
             name = platform.replace("-", "_").lower(),
             os_name = os_name,
             arch_name = arch_name,
             config_settings = config_settings,
-            env = env,
+            want_abis = want_abis or [
+                "cp{0}{1}",
+                "abi3",
+                "none",
+            ],
+            platform_tags = platform_tags,
+            env = {
+                # default to this
+                "implementation_name": "cpython",
+            } | env,
         )
     else:
         config["platforms"].pop(platform)
 
-def _create_config(defaults):
-    if defaults["platforms"]:
-        return struct(**defaults)
+def _set_defaults(defaults):
+    """Set defaults that rules_python is operating under.
+
+    Because this code is also tested in unit tests, leaving it in MODULE.bazel would be
+    a little problematic.
+    """
 
     # NOTE: We have this so that it is easier to maintain unit tests assuming certain
     # defaults
@@ -424,42 +449,62 @@ def _create_config(defaults):
             arch_name = cpu,
             os_name = "linux",
             platform = "linux_{}".format(cpu),
+            want_abis = [],
             config_settings = [
                 "@platforms//os:linux",
                 "@platforms//cpu:{}".format(cpu),
             ],
-            env = {"platform_version": "0"},
+            platform_tags = [
+                "linux_*_{}".format(cpu),
+                "manylinux_*_{}".format(cpu),
+            ],
+            env = {
+                "platform_version": "0",
+            },
         )
-    for cpu in [
-        "aarch64",
-        "x86_64",
-    ]:
+    for cpu, platform_tag_cpus in {
+        "aarch64": ["universal2", "arm64"],
+        "x86_64": ["universal2", "x86_64"],
+    }.items():
         _configure(
             defaults,
             arch_name = cpu,
-            # We choose the oldest non-EOL version at the time when we release `rules_python`.
-            # See https://endoflife.date/macos
             os_name = "osx",
             platform = "osx_{}".format(cpu),
             config_settings = [
                 "@platforms//os:osx",
                 "@platforms//cpu:{}".format(cpu),
             ],
-            env = {"platform_version": "14.0"},
+            want_abis = [],
+            platform_tags = [
+                "macosx_*_{}".format(suffix)
+                for suffix in platform_tag_cpus
+            ],
+            # We choose the oldest non-EOL version at the time when we release `rules_python`.
+            # See https://endoflife.date/macos
+            env = {
+                "platform_version": "14.0",
+            },
         )
 
-    _configure(
-        defaults,
-        arch_name = "x86_64",
-        os_name = "windows",
-        platform = "windows_x86_64",
-        config_settings = [
-            "@platforms//os:windows",
-            "@platforms//cpu:x86_64",
-        ],
-        env = {"platform_version": "0"},
-    )
-    return struct(**defaults)
+    for cpu, platform_tags in {
+        "x86_64": ["win_amd64"],
+    }.items():
+        _configure(
+            defaults,
+            arch_name = cpu,
+            os_name = "windows",
+            platform = "windows_{}".format(cpu),
+            config_settings = [
+                "@platforms//os:windows",
+                "@platforms//cpu:{}".format(cpu),
+            ],
+            want_abis = [],
+            platform_tags = platform_tags,
+            env = {
+                "platform_version": "0",
+            },
+        )
 
 def parse_modules(
         module_ctx,
@@ -515,6 +560,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
         "enable_pipstar": enable_pipstar,
         "platforms": {},
     }
+    _set_defaults(defaults)
     for mod in module_ctx.modules:
         if not (mod.is_root or mod.name == "rules_python"):
             continue
@@ -527,17 +573,18 @@ You cannot use both the additive_build_content and additive_build_content_file a
                 env = tag.env,
                 os_name = tag.os_name,
                 platform = tag.platform,
+                platform_tags = tag.platform_tags,
+                want_abis = tag.want_abis,
                 override = mod.is_root,
                 # TODO @aignas 2025-05-19: add more attr groups:
                 # * for AUTH - the default `netrc` usage could be configured through a common
                 # attribute.
                 # * for index/downloader config. This includes all of those attributes for
                 # overrides, etc. Index overrides per platform could be also used here.
-                # * for whl selection - selecting preferences of which `platform_tag`s we should use
-                # for what. We could also model the `cp313t` freethreaded as separate platforms.
+                # * for whl selection - We could also model the `cp313t` freethreaded as separate platforms.
             )
 
-    config = _create_config(defaults)
+    config = struct(**defaults)
 
     # TODO @aignas 2025-06-03: Merge override API with the builder?
     _overriden_whl_set = {}
@@ -666,7 +713,14 @@ You cannot use both the additive_build_content and additive_build_content_file a
             for whl_name, aliases in out.extra_aliases.items():
                 extra_aliases[hub_name].setdefault(whl_name, {}).update(aliases)
             exposed_packages.setdefault(hub_name, {}).update(out.exposed_packages)
-            whl_libraries.update(out.whl_libraries)
+            for whl_name, lib in out.whl_libraries.items():
+                if enable_pipstar:
+                    whl_libraries.setdefault(whl_name, lib)
+                elif whl_name in lib:
+                    fail("'{}' already in created".format(whl_name))
+                else:
+                    # replicate whl_libraries.update(out.whl_libraries)
+                    whl_libraries[whl_name] = lib
 
             # TODO @aignas 2024-04-05: how do we support different requirement
             # cycles for different abis/oses? For now we will need the users to
@@ -829,25 +883,6 @@ The list of labels to `config_setting` targets that need to be matched for the p
 selected.
 """,
     ),
-    "os_name": attr.string(
-        doc = """\
-The OS name to be used.
-
-:::{note}
-Either this or the appropriate `env` keys should be specified.
-:::
-""",
-    ),
-    "platform": attr.string(
-        doc = """\
-A platform identifier which will be used as the unique identifier within the extension evaluation.
-If you are defining custom platforms in your project and don't want things to clash, use extension
-[isolation] feature.
-
-[isolation]: https://bazel.build/rules/lib/globals/module#use_extension.isolate
-""",
-    ),
-} | {
     "env": attr.string_dict(
         doc = """\
 The values to use for environment markers when evaluating an expression.
@@ -873,6 +908,40 @@ This is only used if the {envvar}`RULES_PYTHON_ENABLE_PIPSTAR` is enabled.
 """,
     ),
     # The values for PEP508 env marker evaluation during the lock file parsing
+    "os_name": attr.string(
+        doc = """\
+The OS name to be used.
+
+:::{note}
+Either this or the appropriate `env` keys should be specified.
+:::
+""",
+    ),
+    "platform": attr.string(
+        doc = """\
+A platform identifier which will be used as the unique identifier within the extension evaluation.
+If you are defining custom platforms in your project and don't want things to clash, use extension
+[isolation] feature.
+
+[isolation]: https://bazel.build/rules/lib/globals/module#use_extension.isolate
+""",
+    ),
+    "platform_tags": attr.string_list(
+        doc = """\
+A list of `platform_tag` matchers so that we can select the best wheel based on the user
+preference. Per platform we will select a single wheel and the last match from this list will
+take precedence.
+
+The items in this list can contain a single `*` character that is equivalent to `.*` regex match.
+""",
+    ),
+    "want_abis": attr.string_list(
+        doc = """\
+A list of ABIs to select wheels for. The values can be either strings or include template
+parameters like `{0}` which will be replaced with python version parts. e.g. `cp{0}{1}` will
+result in `cp313` given the full python version is `3.13.5`.
+""",
+    ),
 }
 
 _SUPPORTED_PEP508_KEYS = [
diff --git a/python/private/pypi/parse_requirements.bzl b/python/private/pypi/parse_requirements.bzl
index 9c610f11d3..6f9ae56fc7 100644
--- a/python/private/pypi/parse_requirements.bzl
+++ b/python/private/pypi/parse_requirements.bzl
@@ -31,13 +31,14 @@ load("//python/private:repo_utils.bzl", "repo_utils")
 load(":index_sources.bzl", "index_sources")
 load(":parse_requirements_txt.bzl", "parse_requirements_txt")
 load(":pep508_requirement.bzl", "requirement")
-load(":whl_target_platforms.bzl", "select_whls")
+load(":select_whl.bzl", "select_whl")
 
 def parse_requirements(
         ctx,
         *,
         requirements_by_platform = {},
         extra_pip_args = [],
+        platforms = {},
         get_index_urls = None,
         evaluate_markers = None,
         extract_url_srcs = True,
@@ -46,6 +47,7 @@ def parse_requirements(
 
     Args:
         ctx: A context that has .read function that would read contents from a label.
+        platforms: The target platform descriptions.
         requirements_by_platform (label_keyed_string_dict): a way to have
             different package versions (or different packages) for different
             os, arch combinations.
@@ -88,7 +90,7 @@ def parse_requirements(
     requirements = {}
     for file, plats in requirements_by_platform.items():
         if logger:
-            logger.debug(lambda: "Using {} for {}".format(file, plats))
+            logger.trace(lambda: "Using {} for {}".format(file, plats))
         contents = ctx.read(file)
 
         # Parse the requirements file directly in starlark to get the information
@@ -161,7 +163,7 @@ def parse_requirements(
     # VCS package references.
     env_marker_target_platforms = evaluate_markers(ctx, reqs_with_env_markers)
     if logger:
-        logger.debug(lambda: "Evaluated env markers from:\n{}\n\nTo:\n{}".format(
+        logger.trace(lambda: "Evaluated env markers from:\n{}\n\nTo:\n{}".format(
             reqs_with_env_markers,
             env_marker_target_platforms,
         ))
@@ -196,6 +198,7 @@ def parse_requirements(
                 name = name,
                 reqs = reqs,
                 index_urls = index_urls,
+                platforms = platforms,
                 env_marker_target_platforms = env_marker_target_platforms,
                 extract_url_srcs = extract_url_srcs,
                 logger = logger,
@@ -203,7 +206,7 @@ def parse_requirements(
         )
         ret.append(item)
         if not item.is_exposed and logger:
-            logger.debug(lambda: "Package '{}' will not be exposed because it is only present on a subset of platforms: {} out of {}".format(
+            logger.trace(lambda: "Package '{}' will not be exposed because it is only present on a subset of platforms: {} out of {}".format(
                 name,
                 sorted(requirement_target_platforms),
                 sorted(requirements),
@@ -219,38 +222,43 @@ def _package_srcs(
         name,
         reqs,
         index_urls,
+        platforms,
         logger,
         env_marker_target_platforms,
         extract_url_srcs):
     """A function to return sources for a particular package."""
     srcs = {}
     for r in sorted(reqs.values(), key = lambda r: r.requirement_line):
-        whls, sdist = _add_dists(
-            requirement = r,
-            index_urls = index_urls.get(name),
-            logger = logger,
-        )
-
         target_platforms = env_marker_target_platforms.get(r.requirement_line, r.target_platforms)
-        target_platforms = sorted(target_platforms)
+        extra_pip_args = tuple(r.extra_pip_args)
 
-        all_dists = [] + whls
-        if sdist:
-            all_dists.append(sdist)
+        for target_platform in target_platforms:
+            if platforms and target_platform not in platforms:
+                fail("The target platform '{}' could not be found in {}".format(
+                    target_platform,
+                    platforms.keys(),
+                ))
 
-        if extract_url_srcs and all_dists:
-            req_line = r.srcs.requirement
-        else:
-            all_dists = [struct(
-                url = "",
-                filename = "",
-                sha256 = "",
-                yanked = False,
-            )]
-            req_line = r.srcs.requirement_line
+            dist = _add_dists(
+                requirement = r,
+                target_platform = platforms.get(target_platform),
+                index_urls = index_urls.get(name),
+                logger = logger,
+            )
+            if logger:
+                logger.debug(lambda: "The whl dist is: {}".format(dist.filename if dist else dist))
+
+            if extract_url_srcs and dist:
+                req_line = r.srcs.requirement
+            else:
+                dist = struct(
+                    url = "",
+                    filename = "",
+                    sha256 = "",
+                    yanked = False,
+                )
+                req_line = r.srcs.requirement_line
 
-        extra_pip_args = tuple(r.extra_pip_args)
-        for dist in all_dists:
             key = (
                 dist.filename,
                 req_line,
@@ -269,9 +277,9 @@ def _package_srcs(
                     yanked = dist.yanked,
                 ),
             )
-            for p in target_platforms:
-                if p not in entry.target_platforms:
-                    entry.target_platforms.append(p)
+
+            if target_platform not in entry.target_platforms:
+                entry.target_platforms.append(target_platform)
 
     return srcs.values()
 
@@ -325,7 +333,7 @@ def host_platform(ctx):
         repo_utils.get_platforms_cpu_name(ctx),
     )
 
-def _add_dists(*, requirement, index_urls, logger = None):
+def _add_dists(*, requirement, index_urls, target_platform, logger = None):
     """Populate dists based on the information from the PyPI index.
 
     This function will modify the given requirements_by_platform data structure.
@@ -333,6 +341,7 @@ def _add_dists(*, requirement, index_urls, logger = None):
     Args:
         requirement: The result of parse_requirements function.
         index_urls: The result of simpleapi_download.
+        target_platform: The target_platform information.
         logger: A logger for printing diagnostic info.
     """
 
@@ -342,7 +351,7 @@ def _add_dists(*, requirement, index_urls, logger = None):
                 logger.debug(lambda: "Could not detect the filename from the URL, falling back to pip: {}".format(
                     requirement.srcs.url,
                 ))
-            return [], None
+            return None
 
         # Handle direct URLs in requirements
         dist = struct(
@@ -353,12 +362,12 @@ def _add_dists(*, requirement, index_urls, logger = None):
         )
 
         if dist.filename.endswith(".whl"):
-            return [dist], None
+            return dist
         else:
-            return [], dist
+            return dist
 
     if not index_urls:
-        return [], None
+        return None
 
     whls = []
     sdist = None
@@ -401,11 +410,16 @@ def _add_dists(*, requirement, index_urls, logger = None):
             for reason, dists in yanked.items()
         ]))
 
-    # Filter out the wheels that are incompatible with the target_platforms.
-    whls = select_whls(
+    if not target_platform:
+        # The pipstar platforms are undefined here, so we cannot do any matching
+        return sdist
+
+    # Select a single wheel that can work on the target_platform
+    return select_whl(
         whls = whls,
-        want_platforms = requirement.target_platforms,
+        python_version = target_platform.env["python_full_version"],
+        implementation_name = target_platform.env["implementation_name"],
+        want_abis = target_platform.want_abis,
+        platforms = target_platform.platform_tags,
         logger = logger,
-    )
-
-    return whls, sdist
+    ) or sdist
diff --git a/python/private/pypi/select_whl.bzl b/python/private/pypi/select_whl.bzl
new file mode 100644
index 0000000000..3d8a90235f
--- /dev/null
+++ b/python/private/pypi/select_whl.bzl
@@ -0,0 +1,130 @@
+"Select a single wheel that fits the parameters of a target platform."
+
+load("//python/private:version.bzl", "version")
+load(":parse_whl_name.bzl", "parse_whl_name")
+
+# Taken from
+# https://packaging.python.org/en/latest/specifications/platform-compatibility-tags/#python-tag
+_PY_TAGS = {
+    # "py": Generic Python (does not require implementation-specific features)
+    "cpython": "cp",
+    "ironpython": "ip",
+    "jython": "jy",
+    "pypy": "pp",
+}
+_PY = "py"
+
+def _get_priority(*, tag, values, allow_wildcard = True):
+    for priority, wp in enumerate(values):
+        head, sep, tail = wp.partition("*")
+        if "*" in tail:
+            fail("only a single '*' can be present in the matcher")
+        if not allow_wildcard and sep:
+            fail("'*' is not allowed in the matcher")
+
+        for p in tag.split("."):
+            if not sep and p == head:
+                return priority
+            elif sep and p.startswith(head) and p.endswith(tail):
+                return priority
+
+    return None
+
+def select_whl(*, whls, python_version, platforms, want_abis, implementation_name = "cpython", limit = 1, logger = None):
+    """Select a whl that is the most suitable for the given platform.
+
+    Args:
+        whls: {type}`list[struct]` a list of candidates which have a `filename`
+            attribute containing the `whl` filename.
+        python_version: {type}`str` the target python version.
+        platforms: {type}`list[str]` the target platform identifiers that may contain
+            a single `*` character.
+        implementation_name: {type}`str` the `implementation_name` from the target_platform env.
+        want_abis: {type}`str` the ABIs that the target_platform is compatible with.
+        limit: {type}`int` number of wheels to return. Defaults to 1.
+        logger: {type}`struct` the logger instance.
+
+    Returns:
+        {type}`list[struct] | struct | None`, a single struct from the `whls` input
+            argument or `None` if a match is not found. If the `limit` is greater than
+            one, then we will return a list.
+    """
+    py_version = version.parse(python_version, strict = True)
+    candidates = {}
+    implementation = _PY_TAGS.get(implementation_name, implementation_name)
+
+    for whl in whls:
+        parsed = parse_whl_name(whl.filename)
+
+        if parsed.python_tag.startswith(_PY):
+            pass
+        elif not parsed.python_tag.startswith(implementation):
+            if logger:
+                logger.debug(lambda: "Discarding the wheel because the implementation '{}' is not compatible with target implementation '{}'".format(
+                    parsed.python_tag,
+                    implementation,
+                ))
+            continue
+
+        if parsed.python_tag == "py2.py3":
+            min_version = "2"
+        else:
+            min_version = parsed.python_tag[len(implementation):]
+
+        if len(min_version) > 1:
+            min_version = "{}.{}".format(min_version[0], min_version[1:])
+
+        min_whl_py_version = version.parse(min_version, strict = True)
+        if not version.is_ge(py_version, min_whl_py_version):
+            if logger:
+                logger.debug(lambda: "Discarding the wheel because the min version supported based on the wheel ABI tag '{}' ({}) is not compatible with the provided target Python version '{}'".format(
+                    parsed.abi_tag,
+                    min_whl_py_version.string,
+                    py_version.string,
+                ))
+            continue
+
+        abi_priority = _get_priority(
+            tag = parsed.abi_tag,
+            values = want_abis,
+            allow_wildcard = False,
+        )
+        if abi_priority == None:
+            if logger:
+                logger.debug(lambda: "The abi '{}' does not match given list: {}".format(
+                    parsed.abi_tag,
+                    want_abis,
+                ))
+            continue
+        platform_priority = _get_priority(
+            tag = parsed.platform_tag,
+            values = platforms,
+        )
+        if platform_priority == None:
+            if logger:
+                logger.debug(lambda: "The platform_tag '{}' does not match given list: {}".format(
+                    parsed.platform_tag,
+                    platforms,
+                ))
+            continue
+
+        key = (
+            # Ensure that we chose the highest compatible version
+            parsed.python_tag.startswith(implementation),
+            platform_priority,
+            # prefer abi_tags in this order
+            version.key(min_whl_py_version),
+            abi_priority,
+        )
+        candidates.setdefault(key, whl)
+
+    if not candidates:
+        return None
+
+    res = [i[1] for i in sorted(candidates.items())]
+    if logger:
+        logger.debug(lambda: "Sorted candidates:\n{}".format(
+            "\n".join([c.filename for c in res]),
+        ))
+
+    return res[-1] if limit == 1 else res[-limit:]
diff --git a/python/private/pypi/whl_target_platforms.bzl b/python/private/pypi/whl_target_platforms.bzl
index 6ea3f120c3..6c3dd5da83 100644
--- a/python/private/pypi/whl_target_platforms.bzl
+++ b/python/private/pypi/whl_target_platforms.bzl
@@ -16,8 +16,6 @@
 A starlark implementation of the wheel platform tag parsing to get the target platform.
 """
 
-load(":parse_whl_name.bzl", "parse_whl_name")
-
 # The order of the dictionaries is to keep definitions with their aliases next to each
 # other
 _CPU_ALIASES = {
@@ -46,136 +44,6 @@ _OS_PREFIXES = {
     "win": "windows",
 }  # buildifier: disable=unsorted-dict-items
 
-def select_whls(*, whls, want_platforms = [], logger = None):
-    """Select a subset of wheels suitable for target platforms from a list.
-
-    Args:
-        whls(list[struct]): A list of candidates which have a `filename`
-            attribute containing the `whl` filename.
-        want_platforms(str): The platforms in "{abi}_{os}_{cpu}" or "{os}_{cpu}" format.
-        logger: A logger for printing diagnostic messages.
-
-    Returns:
-        A filtered list of items from the `whls` arg where `filename` matches
-        the selected criteria. If no match is found, an empty list is returned.
-    """
-    if not whls:
-        return []
-
-    want_abis = {
-        "abi3": None,
-        "none": None,
-    }
-
-    _want_platforms = {}
-    version_limit = None
-
-    for p in want_platforms:
-        if not p.startswith("cp3"):
-            fail("expected all platforms to start with ABI, but got: {}".format(p))
-
-        abi, _, os_cpu = p.partition("_")
-        abi, _, _ = abi.partition(".")
-        _want_platforms[os_cpu] = None
-
-        # TODO @aignas 2025-04-20: add a test
-        _want_platforms["{}_{}".format(abi, os_cpu)] = None
-
-        version_limit_candidate = int(abi[3:])
-        if not version_limit:
-            version_limit = version_limit_candidate
-        if version_limit and version_limit != version_limit_candidate:
-            fail("Only a single python version is supported for now")
-
-        # For some legacy implementations the wheels may target the `cp3xm` ABI
-        _want_platforms["{}m_{}".format(abi, os_cpu)] = None
-        want_abis[abi] = None
-        want_abis[abi + "m"] = None
-
-        # Also add freethreaded wheels if we find them since we started supporting them
-        _want_platforms["{}t_{}".format(abi, os_cpu)] = None
-        want_abis[abi + "t"] = None
-
-    want_platforms = sorted(_want_platforms)
-
-    candidates = {}
-    for whl in whls:
-        parsed = parse_whl_name(whl.filename)
-
-        if logger:
-            logger.trace(lambda: "Deciding whether to use '{}'".format(whl.filename))
-
-        supported_implementations = {}
-        whl_version_min = 0
-        for tag in parsed.python_tag.split("."):
-            supported_implementations[tag[:2]] = None
-
-            if tag.startswith("cp3") or tag.startswith("py3"):
-                version = int(tag[len("..3"):] or 0)
-            else:
-                # In this case it should be eithor "cp2" or "py2" and we will default
-                # to `whl_version_min` = 0
-                continue
-
-            if whl_version_min == 0 or version < whl_version_min:
-                whl_version_min = version
-
-        if not ("cp" in supported_implementations or "py" in supported_implementations):
-            if logger:
-                logger.trace(lambda: "Discarding the whl because the whl does not support CPython, whl supported implementations are: {}".format(supported_implementations))
-            continue
-
-        if want_abis and parsed.abi_tag not in want_abis:
-            # Filter out incompatible ABIs
-            if logger:
-                logger.trace(lambda: "Discarding the whl because the whl abi did not match")
-            continue
-
-        if whl_version_min > version_limit:
-            if logger:
-                logger.trace(lambda: "Discarding the whl because the whl supported python version is too high")
-            continue
-
-        compatible = False
-        if parsed.platform_tag == "any":
-            compatible = True
-        else:
-            for p in whl_target_platforms(parsed.platform_tag, abi_tag = parsed.abi_tag.strip("m") if parsed.abi_tag.startswith("cp") else None):
-                if p.target_platform in want_platforms:
-                    compatible = True
-                    break
-
-        if not compatible:
-            if logger:
-                logger.trace(lambda: "Discarding the whl because the whl does not support the desired platforms: {}".format(want_platforms))
-            continue
-
-        for implementation in supported_implementations:
-            candidates.setdefault(
-                (
-                    parsed.abi_tag,
-                    parsed.platform_tag,
-                ),
-                {},
-            ).setdefault(
-                (
-                    # prefer cp implementation
-                    implementation == "cp",
-                    # prefer higher versions
-                    whl_version_min,
-                    # prefer abi3 over none
-                    parsed.abi_tag != "none",
-                    # prefer cpx abi over abi3
-                    parsed.abi_tag != "abi3",
-                ),
-                [],
-            ).append(whl)
-
-    return [
-        candidates[key][sorted(v)[-1]][-1]
-        for key, v in candidates.items()
-    ]
-
 def whl_target_platforms(platform_tag, abi_tag = ""):
     """Parse the wheel abi and platform tags and return (os, cpu) tuples.
 
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index 51e11342fc..30093f3227 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -77,21 +77,22 @@ def _parse_modules(env, enable_pipstar = 0, **kwargs):
     )
 
 def _default(
+        *,
         arch_name = None,
         config_settings = None,
         os_name = None,
         platform = None,
+        platform_tags = None,
         env = None,
-        whl_limit = None,
-        whl_platforms = None):
+        want_abis = None):
     return struct(
         arch_name = arch_name,
         os_name = os_name,
         platform = platform,
+        platform_tags = platform_tags or [],
         config_settings = config_settings,
         env = env or {},
-        whl_platforms = whl_platforms,
-        whl_limit = whl_limit,
+        want_abis = want_abis or [],
     )
 
 def _parse(
@@ -441,50 +442,60 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
 
     pypi.exposed_packages().contains_exactly({"pypi": ["torch"]})
     pypi.hub_group_map().contains_exactly({"pypi": {}})
-    pypi.hub_whl_map().contains_exactly({"pypi": {
-        "torch": {
-            "pypi_312_torch_cp312_cp312_linux_x86_64_8800deef": [
-                struct(
-                    config_setting = None,
-                    filename = "torch-2.4.1+cpu-cp312-cp312-linux_x86_64.whl",
-                    target_platforms = None,
-                    version = "3.12",
-                ),
-            ],
-            "pypi_312_torch_cp312_cp312_manylinux_2_17_aarch64_36109432": [
-                struct(
-                    config_setting = None,
-                    filename = "torch-2.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
-                    target_platforms = None,
-                    version = "3.12",
-                ),
-            ],
-            "pypi_312_torch_cp312_cp312_win_amd64_3a570e5c": [
-                struct(
-                    config_setting = None,
-                    filename = "torch-2.4.1+cpu-cp312-cp312-win_amd64.whl",
-                    target_platforms = None,
-                    version = "3.12",
-                ),
-            ],
-            "pypi_312_torch_cp312_none_macosx_11_0_arm64_72b484d5": [
-                struct(
-                    config_setting = None,
-                    filename = "torch-2.4.1-cp312-none-macosx_11_0_arm64.whl",
-                    target_platforms = None,
-                    version = "3.12",
-                ),
-            ],
+    pypi.hub_whl_map().contains_exactly({
+        "pypi": {
+            "torch": {
+                "pypi_312_torch_cp312_cp312_linux_x86_64_8800deef": [
+                    whl_config_setting(
+                        target_platforms = ("cp312_linux_x86_64",),
+                        version = "3.12",
+                    ),
+                ],
+                "pypi_312_torch_cp312_cp312_manylinux_2_17_aarch64_36109432": [
+                    whl_config_setting(
+                        version = "3.12",
+                        target_platforms = ("cp312_linux_aarch64",),
+                    ),
+                ],
+                "pypi_312_torch_cp312_cp312_win_amd64_3a570e5c": [
+                    whl_config_setting(
+                        target_platforms = ("cp312_windows_x86_64",),
+                        version = "3.12",
+                    ),
+                ],
+                "pypi_312_torch_cp312_none_macosx_11_0_arm64_72b484d5": [
+                    whl_config_setting(
+                        target_platforms = ("cp312_osx_aarch64",),
+                        version = "3.12",
+                    ),
+                ],
+                # we are falling back onto pip to install/fail on this
+                # TODO @aignas 2025-07-04: maybe we should just exclude the results for
+                # this?
+                "pypi_312_torch_linux_arm_linux_ppc_linux_s390x": [
+                    whl_config_setting(
+                        version = "3.12",
+                        target_platforms = (
+                            "cp312_linux_arm",
+                            "cp312_linux_ppc",
+                            "cp312_linux_s390x",
+                        ),
+                    ),
+                ],
+                # we are falling back onto pip to install/fail on this
+                "pypi_312_torch_osx_x86_64": [
+                    whl_config_setting(
+                        target_platforms = ("cp312_osx_x86_64",),
+                        version = "3.12",
+                    ),
+                ],
+            },
         },
-    }})
+    })
     pypi.whl_libraries().contains_exactly({
         "pypi_312_torch_cp312_cp312_linux_x86_64_8800deef": {
             "dep_template": "@pypi//{name}:{target}",
-            "experimental_target_platforms": [
-                "linux_x86_64",
-                "osx_x86_64",
-                "windows_x86_64",
-            ],
+            "experimental_target_platforms": ["linux_x86_64"],
             "filename": "torch-2.4.1+cpu-cp312-cp312-linux_x86_64.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1+cpu",
@@ -493,13 +504,7 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
         },
         "pypi_312_torch_cp312_cp312_manylinux_2_17_aarch64_36109432": {
             "dep_template": "@pypi//{name}:{target}",
-            "experimental_target_platforms": [
-                "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
-                "osx_aarch64",
-            ],
+            "experimental_target_platforms": ["linux_aarch64"],
             "filename": "torch-2.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1",
@@ -508,11 +513,7 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
         },
         "pypi_312_torch_cp312_cp312_win_amd64_3a570e5c": {
             "dep_template": "@pypi//{name}:{target}",
-            "experimental_target_platforms": [
-                "linux_x86_64",
-                "osx_x86_64",
-                "windows_x86_64",
-            ],
+            "experimental_target_platforms": ["windows_x86_64"],
             "filename": "torch-2.4.1+cpu-cp312-cp312-win_amd64.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1+cpu",
@@ -521,19 +522,23 @@ torch==2.4.1+cpu ; platform_machine == 'x86_64' \
         },
         "pypi_312_torch_cp312_none_macosx_11_0_arm64_72b484d5": {
             "dep_template": "@pypi//{name}:{target}",
-            "experimental_target_platforms": [
-                "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
-                "osx_aarch64",
-            ],
+            "experimental_target_platforms": ["osx_aarch64"],
             "filename": "torch-2.4.1-cp312-none-macosx_11_0_arm64.whl",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1",
             "sha256": "72b484d5b6cec1a735bf3fa5a1c4883d01748698c5e9cfdbeb4ffab7c7987e0d",
             "urls": ["https://torch.index/whl/cpu/torch-2.4.1-cp312-none-macosx_11_0_arm64.whl"],
         },
+        "pypi_312_torch_linux_arm_linux_ppc_linux_s390x": {
+            "dep_template": "@pypi//{name}:{target}",
+            "python_interpreter_target": "unit_test_interpreter_target",
+            "requirement": "torch==2.4.1 --hash=sha256:1495132f30f722af1a091950088baea383fe39903db06b20e6936fd99402803e --hash=sha256:30be2844d0c939161a11073bfbaf645f1c7cb43f62f46cc6e4df1c119fb2a798 --hash=sha256:36109432b10bd7163c9b30ce896f3c2cca1b86b9765f956a1594f0ff43091e2a --hash=sha256:56ad2a760b7a7882725a1eebf5657abbb3b5144eb26bcb47b52059357463c548 --hash=sha256:5fc1d4d7ed265ef853579caf272686d1ed87cebdcd04f2a498f800ffc53dab71 --hash=sha256:72b484d5b6cec1a735bf3fa5a1c4883d01748698c5e9cfdbeb4ffab7c7987e0d --hash=sha256:a38de2803ee6050309aac032676536c3d3b6a9804248537e38e098d0e14817ec --hash=sha256:d36a8ef100f5bff3e9c3cea934b9e0d7ea277cb8210c7152d34a9a6c5830eadd --hash=sha256:ddddbd8b066e743934a4200b3d54267a46db02106876d21cf31f7da7a96f98ea --hash=sha256:fa27b048d32198cda6e9cff0bf768e8683d98743903b7e5d2b1f5098ded1d343",
+        },
+        "pypi_312_torch_osx_x86_64": {
+            "dep_template": "@pypi//{name}:{target}",
+            "python_interpreter_target": "unit_test_interpreter_target",
+            "requirement": "torch==2.4.1+cpu --hash=sha256:0c0a7cc4f7c74ff024d5a5e21230a01289b65346b27a626f6c815d94b4b8c955 --hash=sha256:1dd062d296fb78aa7cfab8690bf03704995a821b5ef69cfc807af5c0831b4202 --hash=sha256:2b03e20f37557d211d14e3fb3f71709325336402db132a1e0dd8b47392185baf --hash=sha256:330e780f478707478f797fdc82c2a96e9b8c5f60b6f1f57bb6ad1dd5b1e7e97e --hash=sha256:3a570e5c553415cdbddfe679207327b3a3806b21c6adea14fba77684d1619e97 --hash=sha256:3c99506980a2fb4b634008ccb758f42dd82f93ae2830c1e41f64536e310bf562 --hash=sha256:76a6fe7b10491b650c630bc9ae328df40f79a948296b41d3b087b29a8a63cbad --hash=sha256:833490a28ac156762ed6adaa7c695879564fa2fd0dc51bcf3fdb2c7b47dc55e6 --hash=sha256:8800deef0026011d502c0c256cc4b67d002347f63c3a38cd8e45f1f445c61364 --hash=sha256:c4f2c3c026e876d4dad7629170ec14fff48c076d6c2ae0e354ab3fdc09024f00",
+        },
     })
     pypi.whl_mods().contains_exactly({})
 
@@ -734,78 +739,99 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
         "pypi": {
             "direct_sdist_without_sha": {
                 "pypi_315_any_name": [
-                    struct(
-                        config_setting = None,
-                        filename = "any-name.tar.gz",
-                        target_platforms = None,
+                    whl_config_setting(
+                        target_platforms = (
+                            "cp315_linux_aarch64",
+                            "cp315_linux_arm",
+                            "cp315_linux_ppc",
+                            "cp315_linux_s390x",
+                            "cp315_linux_x86_64",
+                            "cp315_osx_aarch64",
+                            "cp315_osx_x86_64",
+                            "cp315_windows_x86_64",
+                        ),
                         version = "3.15",
                     ),
                 ],
             },
             "direct_without_sha": {
                 "pypi_315_direct_without_sha_0_0_1_py3_none_any": [
-                    struct(
-                        config_setting = None,
-                        filename = "direct_without_sha-0.0.1-py3-none-any.whl",
-                        target_platforms = None,
+                    whl_config_setting(
+                        target_platforms = (
+                            "cp315_linux_aarch64",
+                            "cp315_linux_arm",
+                            "cp315_linux_ppc",
+                            "cp315_linux_s390x",
+                            "cp315_linux_x86_64",
+                            "cp315_osx_aarch64",
+                            "cp315_osx_x86_64",
+                            "cp315_windows_x86_64",
+                        ),
                         version = "3.15",
                     ),
                 ],
             },
             "git_dep": {
                 "pypi_315_git_dep": [
-                    struct(
-                        config_setting = None,
-                        filename = None,
-                        target_platforms = None,
+                    whl_config_setting(
                         version = "3.15",
                     ),
                 ],
             },
             "pip_fallback": {
                 "pypi_315_pip_fallback": [
-                    struct(
-                        config_setting = None,
-                        filename = None,
-                        target_platforms = None,
+                    whl_config_setting(
                         version = "3.15",
                     ),
                 ],
             },
             "simple": {
                 "pypi_315_simple_py3_none_any_deadb00f": [
-                    struct(
-                        config_setting = None,
-                        filename = "simple-0.0.1-py3-none-any.whl",
-                        target_platforms = None,
-                        version = "3.15",
-                    ),
-                ],
-                "pypi_315_simple_sdist_deadbeef": [
-                    struct(
-                        config_setting = None,
-                        filename = "simple-0.0.1.tar.gz",
-                        target_platforms = None,
+                    whl_config_setting(
+                        target_platforms = (
+                            "cp315_linux_aarch64",
+                            "cp315_linux_arm",
+                            "cp315_linux_ppc",
+                            "cp315_linux_s390x",
+                            "cp315_linux_x86_64",
+                            "cp315_osx_aarch64",
+                            "cp315_osx_x86_64",
+                            "cp315_windows_x86_64",
+                        ),
                         version = "3.15",
                     ),
                 ],
             },
             "some_other_pkg": {
                 "pypi_315_some_py3_none_any_deadb33f": [
-                    struct(
-                        config_setting = None,
-                        filename = "some-other-pkg-0.0.1-py3-none-any.whl",
-                        target_platforms = None,
+                    whl_config_setting(
+                        target_platforms = (
+                            "cp315_linux_aarch64",
+                            "cp315_linux_arm",
+                            "cp315_linux_ppc",
+                            "cp315_linux_s390x",
+                            "cp315_linux_x86_64",
+                            "cp315_osx_aarch64",
+                            "cp315_osx_x86_64",
+                            "cp315_windows_x86_64",
+                        ),
                         version = "3.15",
                     ),
                 ],
             },
             "some_pkg": {
                 "pypi_315_some_pkg_py3_none_any_deadbaaf": [
-                    struct(
-                        config_setting = None,
-                        filename = "some_pkg-0.0.1-py3-none-any.whl",
-                        target_platforms = None,
+                    whl_config_setting(
+                        target_platforms = (
+                            "cp315_linux_aarch64",
+                            "cp315_linux_arm",
+                            "cp315_linux_ppc",
+                            "cp315_linux_s390x",
+                            "cp315_linux_x86_64",
+                            "cp315_osx_aarch64",
+                            "cp315_osx_x86_64",
+                            "cp315_windows_x86_64",
+                        ),
                         version = "3.15",
                     ),
                 ],
@@ -880,25 +906,6 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
             "sha256": "deadb00f",
             "urls": ["example2.org"],
         },
-        "pypi_315_simple_sdist_deadbeef": {
-            "dep_template": "@pypi//{name}:{target}",
-            "experimental_target_platforms": [
-                "linux_aarch64",
-                "linux_arm",
-                "linux_ppc",
-                "linux_s390x",
-                "linux_x86_64",
-                "osx_aarch64",
-                "osx_x86_64",
-                "windows_x86_64",
-            ],
-            "extra_pip_args": ["--extra-args-for-sdist-building"],
-            "filename": "simple-0.0.1.tar.gz",
-            "python_interpreter_target": "unit_test_interpreter_target",
-            "requirement": "simple==0.0.1",
-            "sha256": "deadbeef",
-            "urls": ["example.org"],
-        },
         "pypi_315_some_pkg_py3_none_any_deadbaaf": {
             "dep_template": "@pypi//{name}:{target}",
             "experimental_target_platforms": [
@@ -1060,6 +1067,18 @@ def _test_pipstar_platforms(env):
                         ("linux", "x86_64"),
                         ("osx", "aarch64"),
                     ]
+                ] + [
+                    _default(platform = name)
+                    for name in [
+                        "linux_x86_64",
+                        "linux_aarch64",
+                        "linux_arm",
+                        "linux_ppc",
+                        "linux_s390x",
+                        "osx_x86_64",
+                        "osx_aarch64",
+                        "windows_x86_64",
+                    ]
                 ],
                 parse = [
                     _parse(
diff --git a/tests/pypi/parse_requirements/parse_requirements_tests.bzl b/tests/pypi/parse_requirements/parse_requirements_tests.bzl
index 82fdd0a051..d862d0ebde 100644
--- a/tests/pypi/parse_requirements/parse_requirements_tests.bzl
+++ b/tests/pypi/parse_requirements/parse_requirements_tests.bzl
@@ -16,6 +16,8 @@
 
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("//python/private/pypi:parse_requirements.bzl", "parse_requirements", "select_requirement")  # buildifier: disable=bzl-visibility
+load("//python/private/pypi:pep508_env.bzl", pep508_env = "env")  # buildifier: disable=bzl-visibility
+load("//python/private/pypi:pep508_platform.bzl", "platform_from_str")  # buildifier: disable=bzl-visibility
 
 def _mock_ctx():
     testdata = {
@@ -522,6 +524,18 @@ def _test_overlapping_shas_with_index_results(env):
             "requirements_linux": ["cp39_linux_x86_64"],
             "requirements_osx": ["cp39_osx_x86_64"],
         },
+        platforms = {
+            "cp39_linux_x86_64": struct(
+                platform_tags = ["any"],
+                env = pep508_env(platform_from_str("cp39_linux_x86_64", "")),
+                want_abis = ["none"],
+            ),
+            "cp39_osx_x86_64": struct(
+                platform_tags = ["macosx_*"],
+                env = pep508_env(platform_from_str("cp39_linux_x86_64", "")),
+                want_abis = ["none"],
+            ),
+        },
         get_index_urls = lambda _, __: {
             "foo": struct(
                 sdists = {
@@ -563,20 +577,10 @@ def _test_overlapping_shas_with_index_results(env):
                     filename = "foo-0.0.1-py3-none-any.whl",
                     requirement_line = "foo==0.0.3",
                     sha256 = "deadbaaf",
-                    target_platforms = ["cp39_linux_x86_64", "cp39_osx_x86_64"],
+                    target_platforms = ["cp39_linux_x86_64"],
                     url = "super2",
                     yanked = False,
                 ),
-                struct(
-                    distribution = "foo",
-                    extra_pip_args = [],
-                    filename = "foo-0.0.1.tar.gz",
-                    requirement_line = "foo==0.0.3",
-                    sha256 = "5d15t",
-                    target_platforms = ["cp39_linux_x86_64", "cp39_osx_x86_64"],
-                    url = "sdist",
-                    yanked = False,
-                ),
                 struct(
                     distribution = "foo",
                     extra_pip_args = [],
diff --git a/tests/pypi/select_whl/BUILD.bazel b/tests/pypi/select_whl/BUILD.bazel
new file mode 100644
index 0000000000..0ad8cba0cd
--- /dev/null
+++ b/tests/pypi/select_whl/BUILD.bazel
@@ -0,0 +1,3 @@
+load(":select_whl_tests.bzl", "select_whl_test_suite")
+
+select_whl_test_suite(name = "select_whl_tests")
diff --git a/tests/pypi/whl_target_platforms/select_whl_tests.bzl b/tests/pypi/select_whl/select_whl_tests.bzl
similarity index 57%
rename from tests/pypi/whl_target_platforms/select_whl_tests.bzl
rename to tests/pypi/select_whl/select_whl_tests.bzl
index 1674ac5ef2..57ff39e76a 100644
--- a/tests/pypi/whl_target_platforms/select_whl_tests.bzl
+++ b/tests/pypi/select_whl/select_whl_tests.bzl
@@ -1,22 +1,8 @@
-# Copyright 2024 The Bazel Authors. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 ""
 
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("//python/private:repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "REPO_VERBOSITY_ENV_VAR", "repo_utils")  # buildifier: disable=bzl-visibility
-load("//python/private/pypi:whl_target_platforms.bzl", "select_whls")  # buildifier: disable=bzl-visibility
+load("//python/private/pypi:select_whl.bzl", "select_whl")  # buildifier: disable=bzl-visibility
 
 WHL_LIST = [
     "pkg-0.0.1-cp311-cp311-macosx_10_9_universal2.whl",
@@ -70,6 +56,9 @@ WHL_LIST = [
     "pkg-0.0.1-py310-abi3-any.whl",
     "pkg-0.0.1-py3-abi3-any.whl",
     "pkg-0.0.1-py3-none-any.whl",
+    # Extra examples that should be discarded
+    "pkg-0.0.1-py27-cp27mu-win_amd64.whl",
+    "pkg-0.0.1-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
 ]
 
 def _match(env, got, *want_filenames):
@@ -77,15 +66,16 @@ def _match(env, got, *want_filenames):
         env.expect.that_collection(got).has_size(len(want_filenames))
         return
 
+    got = [g for g in got if g]
     got_filenames = [g.filename for g in got]
-    env.expect.that_collection(got_filenames).contains_exactly(want_filenames)
+    env.expect.that_collection(got_filenames).contains_exactly(want_filenames).in_order()
 
     if got:
         # Check that we pass the original structs
         env.expect.that_str(got[0].other).equals("dummy")
 
-def _select_whls(whls, debug = False, **kwargs):
-    return select_whls(
+def _select_whl(whls, debug = False, **kwargs):
+    return select_whl(
         whls = [
             struct(
                 filename = f,
@@ -107,203 +97,189 @@ def _select_whls(whls, debug = False, **kwargs):
 _tests = []
 
 def _test_simplest(env):
-    got = _select_whls(
-        whls = [
-            "pkg-0.0.1-py2.py3-abi3-any.whl",
-            "pkg-0.0.1-py3-abi3-any.whl",
-            "pkg-0.0.1-py3-none-any.whl",
-        ],
-        want_platforms = ["cp30_ignored"],
+    whls = [
+        "pkg-0.0.1-py2.py3-abi3-any.whl",
+        "pkg-0.0.1-py3-abi3-any.whl",
+        "pkg-0.0.1-py3-none-any.whl",
+    ]
+
+    got = _select_whl(
+        whls = whls,
+        platforms = ["any"],
+        want_abis = ["abi3"],
+        python_version = "3.0",
     )
     _match(
         env,
-        got,
+        [got],
         "pkg-0.0.1-py3-abi3-any.whl",
-        "pkg-0.0.1-py3-none-any.whl",
     )
 
 _tests.append(_test_simplest)
 
 def _test_select_by_supported_py_version(env):
+    whls = [
+        "pkg-0.0.1-py2.py3-abi3-any.whl",
+        "pkg-0.0.1-py3-abi3-any.whl",
+        "pkg-0.0.1-py311-abi3-any.whl",
+    ]
+
     for minor_version, match in {
         8: "pkg-0.0.1-py3-abi3-any.whl",
         11: "pkg-0.0.1-py311-abi3-any.whl",
     }.items():
-        got = _select_whls(
-            whls = [
-                "pkg-0.0.1-py2.py3-abi3-any.whl",
-                "pkg-0.0.1-py3-abi3-any.whl",
-                "pkg-0.0.1-py311-abi3-any.whl",
-            ],
-            want_platforms = ["cp3{}_ignored".format(minor_version)],
+        got = _select_whl(
+            whls = whls,
+            platforms = ["any"],
+            want_abis = ["abi3"],
+            python_version = "3.{}".format(minor_version),
         )
-        _match(env, got, match)
+        _match(env, [got], match)
 
 _tests.append(_test_select_by_supported_py_version)
 
 def _test_select_by_supported_cp_version(env):
+    whls = [
+        "pkg-0.0.1-py2.py3-abi3-any.whl",
+        "pkg-0.0.1-py3-abi3-any.whl",
+        "pkg-0.0.1-py311-abi3-any.whl",
+        "pkg-0.0.1-cp311-abi3-any.whl",
+    ]
+
     for minor_version, match in {
         11: "pkg-0.0.1-cp311-abi3-any.whl",
         8: "pkg-0.0.1-py3-abi3-any.whl",
     }.items():
-        got = _select_whls(
-            whls = [
-                "pkg-0.0.1-py2.py3-abi3-any.whl",
-                "pkg-0.0.1-py3-abi3-any.whl",
-                "pkg-0.0.1-py311-abi3-any.whl",
-                "pkg-0.0.1-cp311-abi3-any.whl",
-            ],
-            want_platforms = ["cp3{}_ignored".format(minor_version)],
+        got = _select_whl(
+            whls = whls,
+            platforms = ["any"],
+            want_abis = ["abi3"],
+            python_version = "3.{}".format(minor_version),
         )
-        _match(env, got, match)
+        _match(env, [got], match)
 
 _tests.append(_test_select_by_supported_cp_version)
 
 def _test_supported_cp_version_manylinux(env):
+    whls = [
+        "pkg-0.0.1-py2.py3-none-manylinux_x86_64.whl",
+        "pkg-0.0.1-py3-none-manylinux_x86_64.whl",
+        "pkg-0.0.1-py311-none-manylinux_x86_64.whl",
+        "pkg-0.0.1-cp311-none-manylinux_x86_64.whl",
+    ]
+
     for minor_version, match in {
         8: "pkg-0.0.1-py3-none-manylinux_x86_64.whl",
         11: "pkg-0.0.1-cp311-none-manylinux_x86_64.whl",
     }.items():
-        got = _select_whls(
-            whls = [
-                "pkg-0.0.1-py2.py3-none-manylinux_x86_64.whl",
-                "pkg-0.0.1-py3-none-manylinux_x86_64.whl",
-                "pkg-0.0.1-py311-none-manylinux_x86_64.whl",
-                "pkg-0.0.1-cp311-none-manylinux_x86_64.whl",
-            ],
-            want_platforms = ["cp3{}_linux_x86_64".format(minor_version)],
+        got = _select_whl(
+            whls = whls,
+            platforms = ["manylinux_x86_64"],
+            want_abis = ["none"],
+            python_version = "3.{}".format(minor_version),
         )
-        _match(env, got, match)
+        _match(env, [got], match)
 
 _tests.append(_test_supported_cp_version_manylinux)
 
 def _test_ignore_unsupported(env):
-    got = _select_whls(
-        whls = [
-            "pkg-0.0.1-xx3-abi3-any.whl",
-        ],
-        want_platforms = ["cp30_ignored"],
+    whls = ["pkg-0.0.1-xx3-abi3-any.whl"]
+    got = _select_whl(
+        whls = whls,
+        platforms = ["any"],
+        want_abis = ["none"],
+        python_version = "3.0",
     )
-    _match(env, got)
+    if got:
+        _match(env, [got], None)
 
 _tests.append(_test_ignore_unsupported)
 
 def _test_match_abi_and_not_py_version(env):
     # Check we match the ABI and not the py version
-    got = _select_whls(whls = WHL_LIST, want_platforms = ["cp37_linux_x86_64"])
+    whls = WHL_LIST
+    platforms = [
+        "musllinux_*_x86_64",
+        "manylinux_*_x86_64",
+    ]
+    got = _select_whl(
+        whls = whls,
+        platforms = platforms,
+        want_abis = ["abi3", "cp37m"],
+        python_version = "3.7",
+    )
     _match(
         env,
-        got,
+        [got],
         "pkg-0.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
-        "pkg-0.0.1-cp37-cp37m-musllinux_1_1_x86_64.whl",
-        "pkg-0.0.1-py3-abi3-any.whl",
-        "pkg-0.0.1-py3-none-any.whl",
     )
 
-_tests.append(_test_match_abi_and_not_py_version)
-
-def _test_select_filename_with_many_tags(env):
-    # Check we can select a filename with many platform tags
-    got = _select_whls(whls = WHL_LIST, want_platforms = ["cp39_linux_x86_32"])
+    got = _select_whl(
+        whls = whls,
+        platforms = platforms[::-1],
+        want_abis = ["abi3", "cp37m"],
+        python_version = "3.7",
+    )
     _match(
         env,
-        got,
-        "pkg-0.0.1-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl",
-        "pkg-0.0.1-cp39-cp39-musllinux_1_1_i686.whl",
-        "pkg-0.0.1-cp39-abi3-any.whl",
-        "pkg-0.0.1-py3-none-any.whl",
+        [got],
+        "pkg-0.0.1-cp37-cp37m-musllinux_1_1_x86_64.whl",
     )
 
-_tests.append(_test_select_filename_with_many_tags)
+_tests.append(_test_match_abi_and_not_py_version)
 
-def _test_osx_prefer_arch_specific(env):
-    # Check that we prefer the specific wheel
-    got = _select_whls(
+def _test_select_filename_with_many_tags(env):
+    # Check we can select a filename with many platform tags
+    got = _select_whl(
         whls = WHL_LIST,
-        want_platforms = ["cp311_osx_x86_64", "cp311_osx_x86_32"],
+        platforms = [
+            "any",
+            "musllinux_*_i686",
+            "manylinux_*_i686",
+        ],
+        want_abis = ["none", "abi3", "cp39"],
+        python_version = "3.9",
+        limit = 5,
     )
     _match(
         env,
         got,
-        "pkg-0.0.1-cp311-cp311-macosx_10_9_universal2.whl",
-        "pkg-0.0.1-cp311-cp311-macosx_10_9_x86_64.whl",
-        "pkg-0.0.1-cp39-abi3-any.whl",
         "pkg-0.0.1-py3-none-any.whl",
-    )
-
-    got = _select_whls(whls = WHL_LIST, want_platforms = ["cp311_osx_aarch64"])
-    _match(
-        env,
-        got,
-        "pkg-0.0.1-cp311-cp311-macosx_10_9_universal2.whl",
-        "pkg-0.0.1-cp311-cp311-macosx_11_0_arm64.whl",
+        "pkg-0.0.1-py3-abi3-any.whl",
         "pkg-0.0.1-cp39-abi3-any.whl",
-        "pkg-0.0.1-py3-none-any.whl",
+        "pkg-0.0.1-cp39-cp39-musllinux_1_1_i686.whl",
+        "pkg-0.0.1-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl",
     )
 
-_tests.append(_test_osx_prefer_arch_specific)
+_tests.append(_test_select_filename_with_many_tags)
 
-def _test_osx_fallback_to_universal2(env):
-    # Check that we can use the universal2 if the arm wheel is not available
-    got = _select_whls(
-        whls = [w for w in WHL_LIST if "arm64" not in w],
-        want_platforms = ["cp311_osx_aarch64"],
+def _test_freethreaded_wheels(env):
+    # Check we prefer platform specific wheels
+    got = _select_whl(
+        whls = WHL_LIST,
+        platforms = [
+            "any",
+            "musllinux_*_x86_64",
+        ],
+        want_abis = ["none", "abi3", "cp313", "cp313t"],
+        python_version = "3.13",
+        limit = 8,
     )
     _match(
         env,
         got,
-        "pkg-0.0.1-cp311-cp311-macosx_10_9_universal2.whl",
-        "pkg-0.0.1-cp39-abi3-any.whl",
+        # The last item has the most priority
         "pkg-0.0.1-py3-none-any.whl",
-    )
-
-_tests.append(_test_osx_fallback_to_universal2)
-
-def _test_prefer_manylinux_wheels(env):
-    # Check we prefer platform specific wheels
-    got = _select_whls(whls = WHL_LIST, want_platforms = ["cp39_linux_x86_64"])
-    _match(
-        env,
-        got,
-        "pkg-0.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl",
-        "pkg-0.0.1-cp39-cp39-musllinux_1_1_x86_64.whl",
+        "pkg-0.0.1-py3-abi3-any.whl",
+        "pkg-0.0.1-py310-abi3-any.whl",
         "pkg-0.0.1-cp39-abi3-any.whl",
-        "pkg-0.0.1-py3-none-any.whl",
-    )
-
-_tests.append(_test_prefer_manylinux_wheels)
-
-def _test_freethreaded_wheels(env):
-    # Check we prefer platform specific wheels
-    got = _select_whls(whls = WHL_LIST, want_platforms = ["cp313_linux_x86_64"])
-    _match(
-        env,
-        got,
-        "pkg-0.0.1-cp313-cp313t-musllinux_1_1_x86_64.whl",
-        "pkg-0.0.1-cp313-cp313-musllinux_1_1_x86_64.whl",
-        "pkg-0.0.1-cp313-abi3-musllinux_1_1_x86_64.whl",
         "pkg-0.0.1-cp313-none-musllinux_1_1_x86_64.whl",
-        "pkg-0.0.1-cp39-abi3-any.whl",
-        "pkg-0.0.1-py3-none-any.whl",
-    )
-
-_tests.append(_test_freethreaded_wheels)
-
-def _test_micro_version_freethreaded(env):
-    # Check we prefer platform specific wheels
-    got = _select_whls(whls = WHL_LIST, want_platforms = ["cp313.3_linux_x86_64"])
-    _match(
-        env,
-        got,
-        "pkg-0.0.1-cp313-cp313t-musllinux_1_1_x86_64.whl",
-        "pkg-0.0.1-cp313-cp313-musllinux_1_1_x86_64.whl",
         "pkg-0.0.1-cp313-abi3-musllinux_1_1_x86_64.whl",
-        "pkg-0.0.1-cp313-none-musllinux_1_1_x86_64.whl",
-        "pkg-0.0.1-cp39-abi3-any.whl",
-        "pkg-0.0.1-py3-none-any.whl",
+        "pkg-0.0.1-cp313-cp313-musllinux_1_1_x86_64.whl",
+        "pkg-0.0.1-cp313-cp313t-musllinux_1_1_x86_64.whl",
     )
 
-_tests.append(_test_micro_version_freethreaded)
+_tests.append(_test_freethreaded_wheels)
 
 def select_whl_test_suite(name):
     """Create the test suite.
diff --git a/tests/pypi/whl_target_platforms/BUILD.bazel b/tests/pypi/whl_target_platforms/BUILD.bazel
index 6c35b08d32..fec25af033 100644
--- a/tests/pypi/whl_target_platforms/BUILD.bazel
+++ b/tests/pypi/whl_target_platforms/BUILD.bazel
@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load(":select_whl_tests.bzl", "select_whl_test_suite")
 load(":whl_target_platforms_tests.bzl", "whl_target_platforms_test_suite")
 
-select_whl_test_suite(name = "select_whl_tests")
-
 whl_target_platforms_test_suite(name = "whl_target_platforms_tests")

From 5e66b5b699122678c8671eeeb04a9d182ffbfdd5 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Sun, 6 Jul 2025 13:03:03 +0900
Subject: [PATCH 4/4] feat: freethreaded support for the builder API

Stacked on #3058

This is a continuation of #3058 where we define freethreaded
platforms. They need to be used only for particular python versions
so I included an extra marker configuration attribute where we
are using pipstar marker evaluation before using the platform.

I think this in general will be a useful tool to configure only
particular platforms for particular python versions

Work towards #2548, since this shows how we can define custom platforms
Work towards #2747

TODO:
- [ ] Fix the remaining expectations in the unit tests. Maybe make the
  tests less brittle and define platforms for unit testing.
---
 MODULE.bazel                                  | 65 ++++++++++++----
 python/private/pypi/extension.bzl             | 78 +++++++++++++------
 python/private/pypi/pip_repository.bzl        |  7 +-
 .../pypi/requirements_files_by_platform.bzl   |  8 +-
 .../resolve_target_platforms.py               |  4 +-
 tests/pypi/extension/extension_tests.bzl      | 31 +++++---
 6 files changed, 141 insertions(+), 52 deletions(-)

diff --git a/MODULE.bazel b/MODULE.bazel
index 9a783eeb0e..05a39b59de 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -70,14 +70,22 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
         config_settings = [
             "@platforms//cpu:{}".format(cpu),
             "@platforms//os:linux",
+            "//python/config_settings:_is_py_freethreaded_{}".format(
+                "yes" if freethreaded else "no",
+            ),
         ],
         env = {"platform_version": "0"},
+        marker = "python_version ~= \"3.13\"" if freethreaded else "",
         os_name = "linux",
-        platform = "linux_{}".format(cpu),
+        platform = "linux_{}{}".format(cpu, freethreaded),
         platform_tags = [
             "linux_*_{}".format(cpu),
             "manylinux_*_{}".format(cpu),
         ],
+        want_abis = [
+            "cp{0}{1}t",
+            "none",
+        ] if freethreaded else [],
     )
     for cpu in [
         "x86_64",
@@ -89,6 +97,10 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
         "ppc",
         "s390x",
     ]
+    for freethreaded in [
+        "",
+        "_freethreaded",
+    ]
 ]
 
 [
@@ -97,16 +109,24 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
         config_settings = [
             "@platforms//cpu:{}".format(cpu),
             "@platforms//os:osx",
+            "//python/config_settings:_is_py_freethreaded_{}".format(
+                "yes" if freethreaded else "no",
+            ),
         ],
         # We choose the oldest non-EOL version at the time when we release `rules_python`.
         # See https://endoflife.date/macos
         env = {"platform_version": "14.0"},
+        marker = "python_version ~= \"3.13\"" if freethreaded else "",
         os_name = "osx",
-        platform = "osx_{}".format(cpu),
+        platform = "osx_{}{}".format(cpu, freethreaded),
         platform_tags = [
             "macosx_*_{}".format(suffix)
             for suffix in platform_tag_cpus
         ],
+        want_abis = [
+            "cp{0}{1}t",
+            "none",
+        ] if freethreaded else [],
     )
     for cpu, platform_tag_cpus in {
         "aarch64": [
@@ -118,19 +138,38 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
             "x86_64",
         ],
     }.items()
+    for freethreaded in [
+        "",
+        "_freethreaded",
+    ]
+]
+
+[
+    pip.default(
+        arch_name = "x86_64",
+        config_settings = [
+            "@platforms//cpu:x86_64",
+            "@platforms//os:windows",
+            "//python/config_settings:_is_py_freethreaded_{}".format(
+                "yes" if freethreaded else "no",
+            ),
+        ],
+        env = {"platform_version": "0"},
+        marker = "python_version ~= \"3.13\"" if freethreaded else "",
+        os_name = "windows",
+        platform = "windows_x86_64{}".format(freethreaded),
+        platform_tags = ["win_amd64"],
+        want_abis = [
+            "cp{0}{1}t",
+            "none",
+        ] if freethreaded else [],
+    )
+    for freethreaded in [
+        "",
+        "_freethreaded",
+    ]
 ]
 
-pip.default(
-    arch_name = "x86_64",
-    config_settings = [
-        "@platforms//cpu:x86_64",
-        "@platforms//os:windows",
-    ],
-    env = {"platform_version": "0"},
-    os_name = "windows",
-    platform = "windows_x86_64",
-    platform_tags = ["win_amd64"],
-)
 pip.parse(
     # NOTE @aignas 2024-10-26: We have an integration test that depends on us
     # being able to build sdists for this hub, so explicitly set this to False.
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index 3e83734e8a..c67257af47 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -30,6 +30,7 @@ load(":hub_repository.bzl", "hub_repository", "whl_config_settings_to_json")
 load(":parse_requirements.bzl", "parse_requirements")
 load(":parse_whl_name.bzl", "parse_whl_name")
 load(":pep508_env.bzl", "env")
+load(":pep508_evaluate.bzl", "evaluate")
 load(":pip_repository_attrs.bzl", "ATTRS")
 load(":requirements_files_by_platform.bzl", "requirements_files_by_platform")
 load(":simpleapi_download.bzl", "simpleapi_download")
@@ -83,8 +84,13 @@ def _platforms(*, python_version, minor_mapping, config):
             os = values.os_name,
             arch = values.arch_name,
         )) | values.env
+
+        if values.marker and not evaluate(values.marker, env = env_):
+            continue
+
         platforms[key] = struct(
             env = env_,
+            tripple = "{}_{}_{}".format(abi, values.os_name, values.arch_name),
             want_abis = [
                 v.format(*python_version.split("."))
                 for v in values.want_abis
@@ -190,17 +196,19 @@ def _create_whl_repos(
         whl_group_mapping = {}
         requirement_cycles = {}
 
+    platforms = _platforms(
+        python_version = pip_attr.python_version,
+        minor_mapping = minor_mapping,
+        config = config,
+    )
+
     if evaluate_markers:
         # This is most likely unit tests
         pass
     elif config.enable_pipstar:
         evaluate_markers = lambda _, requirements: evaluate_markers_star(
             requirements = requirements,
-            platforms = _platforms(
-                python_version = pip_attr.python_version,
-                minor_mapping = minor_mapping,
-                config = config,
-            ),
+            platforms = platforms,
         )
     else:
         # NOTE @aignas 2024-08-02: , we will execute any interpreter that we find either
@@ -219,7 +227,14 @@ def _create_whl_repos(
         # spin up a Python interpreter.
         evaluate_markers = lambda module_ctx, requirements: evaluate_markers_py(
             module_ctx,
-            requirements = requirements,
+            requirements = {
+                k: {
+                    # TODO @aignas 2025-07-06: should we leave this as is?
+                    p: platforms[p].tripple
+                    for p in plats
+                }
+                for k, plats in requirements.items()
+            },
             python_interpreter = pip_attr.python_interpreter,
             python_interpreter_target = python_interpreter_target,
             srcs = pip_attr._evaluate_markers_srcs,
@@ -235,18 +250,14 @@ def _create_whl_repos(
             requirements_osx = pip_attr.requirements_darwin,
             requirements_windows = pip_attr.requirements_windows,
             extra_pip_args = pip_attr.extra_pip_args,
-            platforms = sorted(config.platforms),  # here we only need keys
+            platforms = sorted(platforms),  # here we only need keys
             python_version = full_version(
                 version = pip_attr.python_version,
                 minor_mapping = minor_mapping,
             ),
             logger = logger,
         ),
-        platforms = _platforms(
-            python_version = pip_attr.python_version,
-            minor_mapping = minor_mapping,
-            config = config,
-        ),
+        platforms = platforms,
         extra_pip_args = pip_attr.extra_pip_args,
         get_index_urls = get_index_urls,
         evaluate_markers = evaluate_markers,
@@ -320,6 +331,16 @@ def _create_whl_repos(
                 ))
 
             whl_libraries[repo_name] = repo.args
+            if "experimental_target_platforms" in repo.args:
+                whl_libraries[repo_name] |= {
+                    "experimental_target_platforms": sorted({
+                        # TODO @aignas 2025-07-07: this should be solved in a better way
+                        platforms[candidate].tripple.partition("_")[-1]: None
+                        for p in repo.args["experimental_target_platforms"]
+                        for candidate in platforms
+                        if candidate.endswith(p)
+                    }),
+                }
             whl_map.setdefault(whl.name, {})[repo.config_setting] = repo_name
 
     return struct(
@@ -385,7 +406,7 @@ def _whl_repo(*, src, whl_library_args, is_multiple_versions, download_only, net
         ),
     )
 
-def _configure(config, *, platform, os_name, arch_name, config_settings, env = {}, want_abis, platform_tags, override = False):
+def _configure(config, *, platform, os_name, arch_name, config_settings, env = {}, want_abis, platform_tags, marker, override = False):
     """Set the value in the config if the value is provided"""
     config.setdefault("platforms", {})
     if platform and (os_name or arch_name or config_settings or platform_tags or env):
@@ -406,21 +427,25 @@ def _configure(config, *, platform, os_name, arch_name, config_settings, env = {
             # the lowest priority one needs to be the first one
             platform_tags = ["any"] + platform_tags
 
+        want_abis = want_abis or [
+            "cp{0}{1}",
+            "abi3",
+            "none",
+        ]
+        env = {
+            # default to this
+            "implementation_name": "cpython",
+        } | env
+
         config["platforms"][platform] = struct(
             name = platform.replace("-", "_").lower(),
-            os_name = os_name,
             arch_name = arch_name,
             config_settings = config_settings,
-            want_abis = want_abis or [
-                "cp{0}{1}",
-                "abi3",
-                "none",
-            ],
+            env = env,
+            marker = marker,
+            os_name = os_name,
             platform_tags = platform_tags,
-            env = {
-                # default to this
-                "implementation_name": "cpython",
-            } | env,
+            want_abis = want_abis,
         )
     else:
         config["platforms"].pop(platform)
@@ -491,6 +516,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
                 env = tag.env,
                 os_name = tag.os_name,
                 platform = tag.platform,
+                marker = tag.marker,
                 platform_tags = tag.platform_tags,
                 want_abis = tag.want_abis,
                 override = mod.is_root,
@@ -823,6 +849,12 @@ Supported keys:
 ::::{note}
 This is only used if the {envvar}`RULES_PYTHON_ENABLE_PIPSTAR` is enabled.
 ::::
+""",
+    ),
+    "marker": attr.string(
+        doc = """\
+A marker which will be evaluated to disable the target platform for certain python versions. This
+is especially useful when defining freethreaded platform variants.
 """,
     ),
     # The values for PEP508 env marker evaluation during the lock file parsing
diff --git a/python/private/pypi/pip_repository.bzl b/python/private/pypi/pip_repository.bzl
index e63bd6c3d1..3df56f24ff 100644
--- a/python/private/pypi/pip_repository.bzl
+++ b/python/private/pypi/pip_repository.bzl
@@ -94,7 +94,12 @@ def _pip_repository_impl(rctx):
         extra_pip_args = rctx.attr.extra_pip_args,
         evaluate_markers = lambda rctx, requirements: evaluate_markers_py(
             rctx,
-            requirements = requirements,
+            requirements = {
+                # NOTE @aignas 2025-07-07: because we don't distinguish between
+                # freethreaded and non-freethreaded, it is a 1:1 mapping.
+                req: {p: p for p in plats}
+                for req, plats in requirements.items()
+            },
             python_interpreter = rctx.attr.python_interpreter,
             python_interpreter_target = rctx.attr.python_interpreter_target,
             srcs = rctx.attr._evaluate_markers_srcs,
diff --git a/python/private/pypi/requirements_files_by_platform.bzl b/python/private/pypi/requirements_files_by_platform.bzl
index d8d3651461..356bd4416e 100644
--- a/python/private/pypi/requirements_files_by_platform.bzl
+++ b/python/private/pypi/requirements_files_by_platform.bzl
@@ -37,7 +37,9 @@ def _default_platforms(*, filter, platforms):
         if not prefix:
             return platforms
 
-        match = [p for p in platforms if p.startswith(prefix)]
+        match = [p for p in platforms if p.startswith(prefix) or (
+            p.startswith("cp") and p.partition("_")[-1].startswith(prefix)
+        )]
     else:
         match = [p for p in platforms if filter in p]
 
@@ -140,7 +142,7 @@ def requirements_files_by_platform(
     if logger:
         logger.debug(lambda: "Platforms from pip args: {}".format(platforms_from_args))
 
-    default_platforms = [_platform(p, python_version) for p in platforms]
+    default_platforms = platforms
 
     if platforms_from_args:
         lock_files = [
@@ -252,6 +254,6 @@ def requirements_files_by_platform(
 
     ret = {}
     for plat, file in requirements.items():
-        ret.setdefault(file, []).append(plat)
+        ret.setdefault(file, []).append(_platform(plat, python_version = python_version))
 
     return ret
diff --git a/python/private/pypi/requirements_parser/resolve_target_platforms.py b/python/private/pypi/requirements_parser/resolve_target_platforms.py
index c899a943cc..2204437e41 100755
--- a/python/private/pypi/requirements_parser/resolve_target_platforms.py
+++ b/python/private/pypi/requirements_parser/resolve_target_platforms.py
@@ -50,8 +50,8 @@ def main():
         hashes = prefix + hashes
 
         req = Requirement(entry)
-        for p in target_platforms:
-            (platform,) = Platform.from_string(p)
+        for p, tripple in target_platforms.items():
+            (platform,) = Platform.from_string(tripple)
             if not req.marker or req.marker.evaluate(platform.env_markers("")):
                 response.setdefault(requirement_line, []).append(p)
 
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index ec6b08eb60..836b42cc3b 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -58,20 +58,22 @@ def _mod(*, name, default = [], parse = [], override = [], whl_mods = [], is_roo
             whl_mods = whl_mods,
             default = default or [
                 _default(
-                    platform = "{}_{}".format(os, cpu),
+                    platform = "{}_{}{}".format(os, cpu, freethreaded),
                     os_name = os,
                     arch_name = cpu,
+                    want_abis = ["none", "cp{}{}t"] if freethreaded else [],
                     config_settings = [
                         "@platforms//os:{}".format(os),
                         "@platforms//cpu:{}".format(cpu),
                     ],
                     platform_tags = platform_tags,
                 )
-                for (os, cpu), platform_tags in {
-                    ("linux", "x86_64"): ["linux_*_x86_64", "manylinux_*_x86_64"],
-                    ("linux", "aarch64"): ["linux_*_aarch64", "manylinux_*_aarch64"],
-                    ("osx", "aarch64"): ["macosx_*_arm64"],
-                    ("windows", "aarch64"): ["win_arm64"],
+                for (os, cpu, freethreaded), platform_tags in {
+                    ("linux", "x86_64", ""): ["linux_*_x86_64", "manylinux_*_x86_64"],
+                    ("linux", "x86_64", "_freethreaded"): ["linux_*_x86_64", "manylinux_*_x86_64"],
+                    ("linux", "aarch64", ""): ["linux_*_aarch64", "manylinux_*_aarch64"],
+                    ("osx", "aarch64", ""): ["macosx_*_arm64"],
+                    ("windows", "aarch64", ""): ["win_arm64"],
                 }.items()
             ],
         ),
@@ -101,6 +103,7 @@ def _default(
         platform = None,
         platform_tags = None,
         env = None,
+        marker = None,
         want_abis = None):
     return struct(
         arch_name = arch_name,
@@ -109,6 +112,7 @@ def _default(
         platform_tags = platform_tags or [],
         config_settings = config_settings,
         env = env or {},
+        marker = marker or "",
         want_abis = want_abis or [],
     )
 
@@ -337,10 +341,11 @@ torch==2.4.1 ; platform_machine != 'x86_64' \
                     version = "3.15",
                 ),
             ],
-            "pypi_315_torch_linux_x86_64": [
+            "pypi_315_torch_linux_x86_64_linux_x86_64_freethreaded": [
                 whl_config_setting(
                     target_platforms = [
                         "cp315_linux_x86_64",
+                        "cp315_linux_x86_64_freethreaded",
                     ],
                     version = "3.15",
                 ),
@@ -353,7 +358,7 @@ torch==2.4.1 ; platform_machine != 'x86_64' \
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1 --hash=sha256:deadbeef",
         },
-        "pypi_315_torch_linux_x86_64": {
+        "pypi_315_torch_linux_x86_64_linux_x86_64_freethreaded": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "torch==2.4.1+cpu",
@@ -741,6 +746,7 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
                         target_platforms = (
                             "cp315_linux_aarch64",
                             "cp315_linux_x86_64",
+                            "cp315_linux_x86_64_freethreaded",
                             "cp315_osx_aarch64",
                             "cp315_windows_aarch64",
                         ),
@@ -754,6 +760,7 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
                         target_platforms = (
                             "cp315_linux_aarch64",
                             "cp315_linux_x86_64",
+                            "cp315_linux_x86_64_freethreaded",
                             "cp315_osx_aarch64",
                             "cp315_windows_aarch64",
                         ),
@@ -781,6 +788,7 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
                         target_platforms = (
                             "cp315_linux_aarch64",
                             "cp315_linux_x86_64",
+                            "cp315_linux_x86_64_freethreaded",
                             "cp315_osx_aarch64",
                             "cp315_windows_aarch64",
                         ),
@@ -794,6 +802,7 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
                         target_platforms = (
                             "cp315_linux_aarch64",
                             "cp315_linux_x86_64",
+                            "cp315_linux_x86_64_freethreaded",
                             "cp315_osx_aarch64",
                             "cp315_windows_aarch64",
                         ),
@@ -807,6 +816,7 @@ git_dep @ git+https://git.server/repo/project@deadbeefdeadbeef
                         target_platforms = (
                             "cp315_linux_aarch64",
                             "cp315_linux_x86_64",
+                            "cp315_linux_x86_64_freethreaded",
                             "cp315_osx_aarch64",
                             "cp315_windows_aarch64",
                         ),
@@ -960,12 +970,13 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     pypi.hub_whl_map().contains_exactly({
         "pypi": {
             "optimum": {
-                "pypi_315_optimum_linux_aarch64_linux_x86_64": [
+                "pypi_315_optimum_linux_aarch64_linux_x86_64_linux_x86_64_freethreaded": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
                             "cp315_linux_aarch64",
                             "cp315_linux_x86_64",
+                            "cp315_linux_x86_64_freethreaded",
                         ],
                         config_setting = None,
                         filename = None,
@@ -986,7 +997,7 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     })
 
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_optimum_linux_aarch64_linux_x86_64": {
+        "pypi_315_optimum_linux_aarch64_linux_x86_64_linux_x86_64_freethreaded": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime-gpu]==1.17.1",