wip: add notes fix bugs and address comments

aignas · aignas · commit 818f3e3079eb · 2025-07-07T18:56:45.000+09:00
diff --git a/MODULE.bazel b/MODULE.bazel
@@ -74,7 +74,12 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
         env = {"platform_version": "0"},
         os_name = "linux",
         platform = "linux_{}".format(cpu),
-        platform_tags = [
+        whl_abi_tags = [
+            "cp{major}{minor}",
+            "abi3",
+            "none",
+        ],
+        whl_platform_tags = [
             "linux_*_{}".format(cpu),
             "manylinux_*_{}".format(cpu),
         ],
@@ -103,7 +108,12 @@ pip = use_extension("//python/extensions:pip.bzl", "pip")
         env = {"platform_version": "14.0"},
         os_name = "osx",
         platform = "osx_{}".format(cpu),
-        platform_tags = [
+        whl_abi_tags = [
+            "cp{major}{minor}",
+            "abi3",
+            "none",
+        ],
+        whl_platform_tags = [
             "macosx_*_{}".format(suffix)
             for suffix in platform_tag_cpus
         ],
@@ -129,7 +139,12 @@ pip.default(
     env = {"platform_version": "0"},
     os_name = "windows",
     platform = "windows_x86_64",
-    platform_tags = ["win_amd64"],
+    whl_abi_tags = [
+        "cp{major}{minor}",
+        "abi3",
+        "none",
+    ],
+    whl_platform_tags = ["win_amd64"],
 )
 pip.parse(
     # NOTE @aignas 2024-10-26: We have an integration test that depends on us
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
@@ -68,14 +68,23 @@ def _whl_mods_impl(whl_mods_dict):
 
 def _platforms(*, python_version, minor_mapping, config):
     platforms = {}
-    python_version = full_version(
-        version = python_version,
-        minor_mapping = minor_mapping,
+    python_version = version.parse(
+        full_version(
+            version = python_version,
+            minor_mapping = minor_mapping,
+        ),
+        strict = True,
     )
 
     for platform, values in config.platforms.items():
+        # TODO @aignas 2025-07-07: fix the logic here
         implementation = values.env["implementation_name"][:2].lower()
-        abi = "{}3{}".format(implementation, python_version[2:])
+
+        # TODO @aignas 2025-07-07: move the abi construction somewhere else
+        abi = "{impl}{0}{1}.{2}".format(
+            impl = implementation,
+            *python_version.release
+        )
         key = "{}_{}".format(abi, platform)
 
         env_ = env(struct(
@@ -86,10 +95,13 @@ def _platforms(*, python_version, minor_mapping, config):
         platforms[key] = struct(
             env = env_,
             want_abis = [
-                v.format(*python_version.split("."))
-                for v in values.want_abis
+                v.format(
+                    major = python_version.release[0],
+                    minor = python_version.release[1],
+                )
+                for v in values.whl_abi_tags
             ],
-            platform_tags = values.platform_tags,
+            platform_tags = values.whl_platform_tags,
         )
     return platforms
 
@@ -385,10 +397,10 @@ 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 = {}, 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 platform_tags or env):
+    if platform and (os_name or arch_name or config_settings or whl_platform_tags or env):
         if not override and config.get("platforms", {}).get(platform):
             return
 
@@ -402,21 +414,21 @@ 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:
+        if whl_platform_tags and "any" not in whl_platform_tags:
             # the lowest priority one needs to be the first one
-            platform_tags = ["any"] + platform_tags
+            whl_platform_tags = ["any"] + whl_platform_tags
 
         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}",
+            whl_abi_tags = whl_abi_tags or [
+                "cp{major}{minor}",
                 "abi3",
                 "none",
             ],
-            platform_tags = platform_tags,
+            whl_platform_tags = whl_platform_tags,
             env = {
                 # default to this
                 "implementation_name": "cpython",
@@ -491,8 +503,8 @@ 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,
+                whl_abi_tags = tag.whl_abi_tags,
+                whl_platform_tags = tag.whl_platform_tags,
                 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
@@ -844,20 +856,24 @@ If you are defining custom platforms in your project and don't want things to cl
 [isolation]: https://bazel.build/rules/lib/globals/module#use_extension.isolate
 """,
     ),
-    "platform_tags": attr.string_list(
+    "whl_abi_tags": 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`.
+
+See official [docs](https://packaging.python.org/en/latest/specifications/platform-compatibility-tags/#abi-tag) for more information.
+""",
+    ),
+    "whl_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`.
+
+See official [docs](https://packaging.python.org/en/latest/specifications/platform-compatibility-tags/#platform-tag) for more information.
 """,
     ),
 }
diff --git a/python/private/pypi/pep508_env.bzl b/python/private/pypi/pep508_env.bzl
@@ -15,6 +15,8 @@
 """This module is for implementing PEP508 environment definition.
 """
 
+_DEFAULT = "//conditions:default"
+
 # See https://stackoverflow.com/a/45125525
 platform_machine_aliases = {
     # These pairs mean the same hardware, but different values may be used
@@ -59,11 +61,11 @@ platform_machine_select_map = {
     "@platforms//cpu:x86_64": "x86_64",
     # The value is empty string if it cannot be determined:
     # https://docs.python.org/3/library/platform.html#platform.machine
-    "//conditions:default": "",
+    _DEFAULT: "",
 }
 
 # Platform system returns results from the `uname` call.
-_platform_system_values = {
+platform_system_select_map = {
     # See https://peps.python.org/pep-0738/#platform
     "android": "Android",
     "freebsd": "FreeBSD",
@@ -75,15 +77,9 @@ _platform_system_values = {
     "openbsd": "OpenBSD",
     "osx": "Darwin",
     "windows": "Windows",
-}
-
-platform_system_select_map = {
-    "@platforms//os:{}".format(bazel_os): py_system
-    for bazel_os, py_system in _platform_system_values.items()
-} | {
     # The value is empty string if it cannot be determined:
     # https://docs.python.org/3/library/platform.html#platform.machine
-    "//conditions:default": "",
+    _DEFAULT: "",
 }
 
 # The copy of SO [answer](https://stackoverflow.com/a/13874620) containing
@@ -111,50 +107,43 @@ platform_system_select_map = {
 # (**) Prior Python 3.8 could also be aix5 or aix7; use sys.platform.startswith()
 #
 # We are using only the subset that we actually support.
-_sys_platform_values = {
+sys_platform_select_map = {
     # These values are decided by the sys.platform docs.
-    "android": "android",
-    "emscripten": "emscripten",
+    "@platforms//os:android": "android",
+    "@platforms//os:emscripten": "emscripten",
     # NOTE: The below values are approximations. The sys.platform() docs
     # don't have documented values for these OSes. Per docs, the
     # sys.platform() value reflects the OS at the time Python was *built*
     # instead of the runtime (target) OS value.
-    "freebsd": "freebsd",
-    "ios": "ios",
-    "linux": "linux",
-    "openbsd": "openbsd",
-    "osx": "darwin",
-    "wasi": "wasi",
-    "windows": "win32",
-}
-
-sys_platform_select_map = {
-    "@platforms//os:{}".format(bazel_os): py_platform
-    for bazel_os, py_platform in _sys_platform_values.items()
-} | {
+    "@platforms//os:freebsd": "freebsd",
+    "@platforms//os:ios": "ios",
+    "@platforms//os:linux": "linux",
+    "@platforms//os:openbsd": "openbsd",
+    "@platforms//os:osx": "darwin",
+    "@platforms//os:wasi": "wasi",
+    "@platforms//os:windows": "win32",
     # For lack of a better option, use empty string. No standard doc/spec
     # about sys_platform value.
-    "//conditions:default": "",
+    _DEFAULT: "",
 }
 
 # The "java" value is documented, but with Jython defunct,
 # shouldn't occur in practice.
 # The os.name value is technically a property of the runtime, not the
 # targetted runtime OS, but the distinction shouldn't matter if
 # things are properly configured.
-_os_name_values = {
-    "linux": "posix",
-    "osx": "posix",
-    "windows": "nt",
-}
-
 os_name_select_map = {
-    "@platforms//os:{}".format(bazel_os): py_os
-    for bazel_os, py_os in _os_name_values.items()
-} | {
-    "//conditions:default": "posix",
+    "@platforms//os:windows": "nt",
+    _DEFAULT: "posix",
 }
 
+# TODO @aignas 2025-07-07: add a test to ensure that this is tested
+def _get_from_map(m, key):
+    if _DEFAULT in m:
+        return m.get(key, m[_DEFAULT])
+    else:
+        return m[key]
+
 def env(target_platform, *, extra = None):
     """Return an env target platform
 
@@ -174,20 +163,23 @@ def env(target_platform, *, extra = None):
         env["extra"] = extra
 
     if target_platform.abi:
+        # TODO @aignas 2025-07-07: do not parse the version like this here
         minor_version, _, micro_version = target_platform.abi[3:].partition(".")
         micro_version = micro_version or "0"
         env = env | {
             "implementation_version": "3.{}.{}".format(minor_version, micro_version),
             "python_full_version": "3.{}.{}".format(minor_version, micro_version),
             "python_version": "3.{}".format(minor_version),
         }
+
     if target_platform.os and target_platform.arch:
-        os = target_platform.os
+        os = "@platforms//os:{}".format(target_platform.os)
+        arch = "@platforms//cpu:{}".format(target_platform.arch)
         env = env | {
-            "os_name": _os_name_values.get(os, ""),
-            "platform_machine": target_platform.arch,
-            "platform_system": _platform_system_values.get(os, ""),
-            "sys_platform": _sys_platform_values.get(os, ""),
+            "os_name": _get_from_map(os_name_select_map, os),
+            "platform_machine": _get_from_map(platform_machine_select_map, arch),
+            "platform_system": _get_from_map(platform_system_select_map, os),
+            "sys_platform": _get_from_map(sys_platform_select_map, os),
         }
     set_missing_env_defaults(env)
 
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
@@ -65,9 +65,9 @@ def _mod(*, name, default = [], parse = [], override = [], whl_mods = [], is_roo
                         "@platforms//os:{}".format(os),
                         "@platforms//cpu:{}".format(cpu),
                     ],
-                    platform_tags = platform_tags,
+                    whl_platform_tags = whl_platform_tags,
                 )
-                for (os, cpu), platform_tags in {
+                for (os, cpu), whl_platform_tags in {
                     ("linux", "x86_64"): ["linux_*_x86_64", "manylinux_*_x86_64"],
                     ("linux", "aarch64"): ["linux_*_aarch64", "manylinux_*_aarch64"],
                     ("osx", "aarch64"): ["macosx_*_arm64"],
@@ -99,17 +99,17 @@ def _default(
         config_settings = None,
         os_name = None,
         platform = None,
-        platform_tags = None,
+        whl_platform_tags = None,
         env = None,
-        want_abis = None):
+        whl_abi_tags = None):
     return struct(
         arch_name = arch_name,
         os_name = os_name,
         platform = platform,
-        platform_tags = platform_tags or [],
+        whl_platform_tags = whl_platform_tags or [],
         config_settings = config_settings,
         env = env or {},
-        want_abis = want_abis or [],
+        whl_abi_tags = whl_abi_tags or [],
     )
 
 def _parse(
@@ -407,9 +407,9 @@ def _test_torch_experimental_index_url(env):
                             "@platforms//os:{}".format(os),
                             "@platforms//cpu:{}".format(cpu),
                         ],
-                        platform_tags = platform_tags,
+                        whl_platform_tags = whl_platform_tags,
                     )
-                    for (os, cpu), platform_tags in {
+                    for (os, cpu), whl_platform_tags in {
                         ("linux", "x86_64"): ["linux_*_x86_64", "manylinux_*_x86_64"],
                         ("linux", "aarch64"): ["linux_*_aarch64", "manylinux_*_aarch64"],
                         ("osx", "aarch64"): ["macosx_*_arm64"],
