feat(pypi): freethreaded support for the builder API

aignas · aignas · commit 78d4007b4fe1 · 2025-07-21T13:52:34.000+09:00
DO NOT MERGE: stacked on bazel-contrib#3058 This is a continuation of bazel-contrib#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 bazel-contrib#2548, since this shows how we can define custom platforms Work towards bazel-contrib#2747
diff --git a/MODULE.bazel b/MODULE.bazel
@@ -70,11 +70,15 @@ 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),
-        whl_abi_tags = [
+        platform = "linux_{}{}".format(cpu, freethreaded),
+        whl_abi_tags = ["cp{major}{minor}t"] if freethreaded else [
             "abi3",
             "cp{major}{minor}",
         ],
@@ -87,6 +91,10 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
         "x86_64",
         "aarch64",
     ]
+    for freethreaded in [
+        "",
+        "_freethreaded",
+    ]
 ]
 
 [
@@ -95,13 +103,17 @@ 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),
-        whl_abi_tags = [
+        platform = "osx_{}{}".format(cpu, freethreaded),
+        whl_abi_tags = ["cp{major}{minor}t"] if freethreaded else [
             "abi3",
             "cp{major}{minor}",
         ],
@@ -120,6 +132,10 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
             "x86_64",
         ],
     }.items()
+    for freethreaded in [
+        "",
+        "_freethreaded",
+    ]
 ]
 
 [
@@ -128,11 +144,15 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
         config_settings = [
             "@platforms//cpu:{}".format(cpu),
             "@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_{}".format(cpu),
-        whl_abi_tags = [
+        platform = "windows_{}{}".format(cpu, freethreaded),
+        whl_abi_tags = ["cp{major}{minor}t"] if freethreaded else [
             "abi3",
             "cp{major}{minor}",
         ],
@@ -141,6 +161,10 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
     for cpu, whl_platform_tags in {
         "x86_64": ["win_amd64"],
     }.items()
+    for freethreaded in [
+        "",
+        "_freethreaded",
+    ]
 ]
 
 pip.parse(
diff --git 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(":python_tag.bzl", "python_tag")
 load(":requirements_files_by_platform.bzl", "requirements_files_by_platform")
@@ -80,21 +81,27 @@ def _platforms(*, python_version, minor_mapping, config):
     for platform, values in config.platforms.items():
         # TODO @aignas 2025-07-07: this is probably doing the parsing of the version too
         # many times.
-        key = "{}{}{}.{}_{}".format(
+        abi = "{}{}{}.{}".format(
             python_tag(values.env["implementation_name"]),
             python_version.release[0],
             python_version.release[1],
             python_version.release[2],
-            platform,
         )
+        key = "{}_{}".format(abi, platform)
+
+        env_ = env(
+            env = values.env,
+            os = values.os_name,
+            arch = values.arch_name,
+            python_version = python_version.string,
+        )
+
+        if values.marker and not evaluate(values.marker, env = env_):
+            continue
 
         platforms[key] = struct(
-            env = env(
-                env = values.env,
-                os = values.os_name,
-                arch = values.arch_name,
-                python_version = python_version.string,
-            ),
+            env = env_,
+            triple = "{}_{}_{}".format(abi, values.os_name, values.arch_name),
             whl_abi_tags = [
                 v.format(
                     major = python_version.release[0],
@@ -203,17 +210,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
@@ -232,7 +241,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].triple
+                    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,
@@ -248,18 +264,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,
@@ -346,6 +358,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].triple.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(
@@ -434,14 +456,15 @@ def _configure(
         arch_name,
         config_settings,
         env = {},
+        marker,
         whl_abi_tags,
         whl_platform_tags,
         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 whl_abi_tags or whl_platform_tags or env
+        os_name or arch_name or config_settings or whl_abi_tags or whl_platform_tags or env or marker
     ):
         if not override and config["platforms"].get(platform):
             return
@@ -455,6 +478,7 @@ def _configure(
             "arch_name": arch_name,
             "config_settings": config_settings,
             "env": env,
+            "marker": marker,
             "name": platform.replace("-", "_").lower(),
             "os_name": os_name,
             "whl_abi_tags": whl_abi_tags,
@@ -470,7 +494,7 @@ def _configure(
     else:
         config["platforms"].pop(platform)
 
-def _plat(*, name, arch_name, os_name, config_settings = [], env = {}, whl_abi_tags = [], whl_platform_tags = []):
+def _plat(*, name, arch_name, os_name, config_settings = [], env = {}, marker = "", whl_abi_tags = [], whl_platform_tags = []):
     # NOTE @aignas 2025-07-08: the least preferred is the first item in the list
     if "any" not in whl_platform_tags:
         # the lowest priority one needs to be the first one
@@ -490,6 +514,7 @@ def _plat(*, name, arch_name, os_name, config_settings = [], env = {}, whl_abi_t
             # defaults for env
             "implementation_name": "cpython",
         } | env,
+        marker = marker,
         whl_abi_tags = whl_abi_tags,
         whl_platform_tags = whl_platform_tags,
     )
@@ -524,6 +549,7 @@ def build_config(
                 config_settings = tag.config_settings,
                 env = tag.env,
                 os_name = tag.os_name,
+                marker = tag.marker,
                 platform = platform,
                 override = mod.is_root,
                 whl_abi_tags = tag.whl_abi_tags,
@@ -533,8 +559,6 @@ def build_config(
                 # 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.
             )
 
     return struct(
@@ -918,6 +942,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
@@ -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
@@ -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
@@ -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
