feat(pypi): incrementally build platform configuration

aignas · aignas · commit f2e68cf93027 · 2025-07-21T13:08:05.000+09:00
Before this PR the configuration for platforms would be built non-incrementally, making it harder for users to override particular attributes of the already configured ones. With this PR the new features introduced in bazel-contrib#3058 will be easier to override. Work towards bazel-contrib#2747
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
@@ -380,24 +380,90 @@ def _whl_repo(*, src, whl_library_args, is_multiple_versions, download_only, net
 def _configure(config, *, platform, os_name, arch_name, config_settings, env = {}, 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 env):
         if not override and config.get("platforms", {}).get(platform):
             return
 
         for key in env:
             if key not in _SUPPORTED_PEP508_KEYS:
                 fail("Unsupported key in the PEP508 environment: {}".format(key))
 
-        config["platforms"][platform] = struct(
-            name = platform.replace("-", "_").lower(),
-            os_name = os_name,
-            arch_name = arch_name,
-            config_settings = config_settings,
-            env = env,
-        )
+        config["platforms"].setdefault(platform, {})
+        for key, value in {
+            "arch_name": arch_name,
+            "config_settings": config_settings,
+            "env": env,
+            "name": platform.replace("-", "_").lower(),
+            "os_name": os_name,
+        }.items():
+            if not value:
+                continue
+
+            if not override and config.get(key):
+                continue
+
+            config["platforms"][platform][key] = value
     else:
         config["platforms"].pop(platform)
 
+def _plat(*, name, arch_name, os_name, config_settings = [], env = {}):
+    return struct(
+        name = name,
+        arch_name = arch_name,
+        os_name = os_name,
+        config_settings = config_settings,
+        env = env,
+    )
+
+def build_config(
+        *,
+        module_ctx,
+        enable_pipstar):
+    """Parse 'configure' and 'default' extension tags
+
+    Args:
+        module_ctx: {type}`module_ctx` module context.
+        enable_pipstar: {type}`bool` a flag to enable dropping Python dependency for
+            evaluation of the extension.
+
+    Returns:
+        A struct with the configuration.
+    """
+    defaults = {
+        "platforms": {},
+    }
+    for mod in module_ctx.modules:
+        if not (mod.is_root or mod.name == "rules_python"):
+            continue
+
+        platform = None
+        for tag in mod.tags.default:
+            platform = tag.platform or platform
+            _configure(
+                defaults,
+                arch_name = tag.arch_name,
+                config_settings = tag.config_settings,
+                env = tag.env,
+                os_name = tag.os_name,
+                platform = platform,
+                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.
+            )
+
+    return struct(
+        platforms = {
+            name: _plat(**values)
+            for name, values in defaults["platforms"].items()
+        },
+        enable_pipstar = enable_pipstar,
+    )
+
 def parse_modules(
         module_ctx,
         _fail = fail,
@@ -448,33 +514,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
                 srcs_exclude_glob = whl_mod.srcs_exclude_glob,
             )
 
-    defaults = {
-        "enable_pipstar": enable_pipstar,
-        "platforms": {},
-    }
-    for mod in module_ctx.modules:
-        if not (mod.is_root or mod.name == "rules_python"):
-            continue
-
-        for tag in mod.tags.default:
-            _configure(
-                defaults,
-                arch_name = tag.arch_name,
-                config_settings = tag.config_settings,
-                env = tag.env,
-                os_name = tag.os_name,
-                platform = tag.platform,
-                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.
-            )
-
-    config = struct(**defaults)
+    config = build_config(module_ctx = module_ctx, enable_pipstar = enable_pipstar)
 
     # TODO @aignas 2025-06-03: Merge override API with the builder?
     _overriden_whl_set = {}
@@ -651,6 +691,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
             k: dict(sorted(args.items()))
             for k, args in sorted(whl_libraries.items())
         },
+        config = config,
     )
 
 def _pip_impl(module_ctx):
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
@@ -16,7 +16,7 @@
 
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
 load("@rules_testing//lib:truth.bzl", "subjects")
-load("//python/private/pypi:extension.bzl", "parse_modules")  # buildifier: disable=bzl-visibility
+load("//python/private/pypi:extension.bzl", "build_config", "parse_modules")  # buildifier: disable=bzl-visibility
 load("//python/private/pypi:parse_simpleapi_html.bzl", "parse_simpleapi_html")  # buildifier: disable=bzl-visibility
 load("//python/private/pypi:whl_config_setting.bzl", "whl_config_setting")  # buildifier: disable=bzl-visibility
 
@@ -92,6 +92,18 @@ def _parse_modules(env, enable_pipstar = 0, **kwargs):
         ),
     )
 
+def _build_config(env, enable_pipstar = 0, **kwargs):
+    return env.expect.that_struct(
+        build_config(
+            enable_pipstar = enable_pipstar,
+            **kwargs
+        ),
+        attrs = dict(
+            platforms = subjects.dict,
+            enable_pipstar = subjects.bool,
+        ),
+    )
+
 def _default(
         arch_name = None,
         config_settings = None,
@@ -1108,6 +1120,45 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
 
 _tests.append(_test_pipstar_platforms)
 
+def _test_build_pipstar_platform(env):
+    config = _build_config(
+        env,
+        module_ctx = _mock_mctx(
+            _mod(
+                name = "rules_python",
+                default = [
+                    _default(
+                        platform = "myplat",
+                        os_name = "linux",
+                        arch_name = "x86_64",
+                    ),
+                    _default(
+                        config_settings = [
+                            "@platforms//os:linux",
+                            "@platforms//cpu:x86_64",
+                        ],
+                    ),
+                ],
+            ),
+        ),
+        enable_pipstar = True,
+    )
+    config.enable_pipstar().equals(True)
+    config.platforms().contains_exactly({
+        "myplat": struct(
+            name = "myplat",
+            os_name = "linux",
+            arch_name = "x86_64",
+            config_settings = [
+                "@platforms//os:linux",
+                "@platforms//cpu:x86_64",
+            ],
+            env = {},
+        ),
+    })
+
+_tests.append(_test_build_pipstar_platform)
+
 def extension_test_suite(name):
     """Create the test suite.
 
