fix(pypi): reuse select dicts for constructing the env

aignas · aignas · commit 0af9537a975d · 2025-07-21T10:51:19.000+09:00
Before this PR we would be constructing slightly different environments when the `env_marker_setting` is doing it in the analysis phase and when we are doing it in the repo phase due to how the defaults are handled. In this change we simply reuse the same select statements and add an extra helper that is allowing us to process that. Work towards bazel-contrib#2949 Prep for bazel-contrib#3058
diff --git a/python/private/pypi/BUILD.bazel b/python/private/pypi/BUILD.bazel
@@ -252,6 +252,9 @@ bzl_library(
 bzl_library(
     name = "pep508_env_bzl",
     srcs = ["pep508_env.bzl"],
+    deps = [
+        "//python/private:version_bzl",
+    ],
 )
 
 bzl_library(
@@ -263,11 +266,6 @@ bzl_library(
     ],
 )
 
-bzl_library(
-    name = "pep508_platform_bzl",
-    srcs = ["pep508_platform.bzl"],
-)
-
 bzl_library(
     name = "pep508_requirement_bzl",
     srcs = ["pep508_requirement.bzl"],
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
@@ -76,11 +76,12 @@ def _platforms(*, python_version, minor_mapping, config):
 
     for platform, values in config.platforms.items():
         key = "{}_{}".format(abi, platform)
-        platforms[key] = env(struct(
-            abi = abi,
+        platforms[key] = env(
+            env = values.env,
             os = values.os_name,
             arch = values.arch_name,
-        )) | values.env
+            python_version = python_version,
+        )
     return platforms
 
 def _create_whl_repos(
diff --git a/python/private/pypi/pep508_env.bzl b/python/private/pypi/pep508_env.bzl
@@ -15,6 +15,20 @@
 """This module is for implementing PEP508 environment definition.
 """
 
+load("//python/private:version.bzl", "version")
+
+_DEFAULT = "//conditions:default"
+
+# Here we store the aliases in the platform so that the users can specify any valid target in
+# there.
+_cpu_aliases = {
+    "arm": "aarch32",
+    "arm64": "aarch64",
+}
+_os_aliases = {
+    "macos": "osx",
+}
+
 # See https://stackoverflow.com/a/45125525
 platform_machine_aliases = {
     # These pairs mean the same hardware, but different values may be used
@@ -59,7 +73,7 @@ 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.
@@ -73,17 +87,17 @@ _platform_system_values = {
     "linux": "Linux",
     "netbsd": "NetBSD",
     "openbsd": "OpenBSD",
-    "osx": "Darwin",
+    "osx": "Darwin",  # NOTE: macos is an alias to osx, we handle it through _os_aliases
     "windows": "Windows",
 }
 
 platform_system_select_map = {
-    "@platforms//os:{}".format(bazel_os): py_system
-    for bazel_os, py_system in _platform_system_values.items()
+    "@platforms//os:{}".format(os): value
+    for os, value 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
@@ -123,72 +137,78 @@ _sys_platform_values = {
     "ios": "ios",
     "linux": "linux",
     "openbsd": "openbsd",
-    "osx": "darwin",
+    "osx": "darwin",  # NOTE: macos is an alias to osx, we handle it through _os_aliases
     "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()
+    # These values are decided by the sys.platform docs.
+    "@platforms//os:{}".format(os): value
+    for os, value in _sys_platform_values.items()
 } | {
     # 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",
 }
 
-def env(target_platform, *, extra = None):
+def _set_default(env, env_key, m, key):
+    """Set the default value in the env if it is not already set."""
+    default = m.get(key, m[_DEFAULT])
+    env.setdefault(env_key, default)
+
+def env(*, env = None, os, arch, python_version = "", extra = None):
     """Return an env target platform
 
     NOTE: This is for use during the loading phase. For the analysis phase,
     `env_marker_setting()` constructs the env dict.
 
     Args:
-        target_platform: {type}`str` the target platform identifier, e.g.
-            `cp33_linux_aarch64`
+        env: {type}`str` the environment.
+        os: {type}`str` the OS name.
+        arch: {type}`str` the CPU name.
+        python_version: {type}`str` the full python version.
         extra: {type}`str` the extra value to be added into the env.
 
     Returns:
         A dict that can be used as `env` in the marker evaluation.
     """
-    env = create_env()
+    env = env or {}
+    env = env | create_env()
     if extra != None:
         env["extra"] = extra
 
-    if target_platform.abi:
-        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
+    if python_version:
+        v = version.parse(python_version)
+        major = v.release[0]
+        minor = v.release[1]
+        micro = v.release[2] if len(v.release) > 2 else 0
         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, ""),
+            "implementation_version": "{}.{}.{}".format(major, minor, micro),
+            "python_full_version": "{}.{}.{}".format(major, minor, micro),
+            "python_version": "{}.{}".format(major, minor),
         }
+
+    if os:
+        os = "@platforms//os:{}".format(_os_aliases.get(os, os))
+        _set_default(env, "os_name", os_name_select_map, os)
+        _set_default(env, "platform_system", platform_system_select_map, os)
+        _set_default(env, "sys_platform", sys_platform_select_map, os)
+
+    if arch:
+        arch = "@platforms//cpu:{}".format(_cpu_aliases.get(arch, arch))
+        _set_default(env, "platform_machine", platform_machine_select_map, arch)
+
     set_missing_env_defaults(env)
 
     return env
diff --git a/python/private/pypi/pep508_platform.bzl b/python/private/pypi/pep508_platform.bzl
diff --git a/tests/pypi/pep508/BUILD.bazel b/tests/pypi/pep508/BUILD.bazel
@@ -1,11 +1,16 @@
 load(":deps_tests.bzl", "deps_test_suite")
+load(":env_tests.bzl", "env_test_suite")
 load(":evaluate_tests.bzl", "evaluate_test_suite")
 load(":requirement_tests.bzl", "requirement_test_suite")
 
 deps_test_suite(
     name = "deps_tests",
 )
 
+env_test_suite(
+    name = "env_tests",
+)
+
 evaluate_test_suite(
     name = "evaluate_tests",
 )
diff --git a/tests/pypi/pep508/env_tests.bzl b/tests/pypi/pep508/env_tests.bzl
@@ -0,0 +1,69 @@
+"""Tests to check for env construction."""
+
+load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("//python/private/pypi:pep508_env.bzl", pep508_env = "env")  # buildifier: disable=bzl-visibility
+
+_tests = []
+
+def _test_env_defaults(env):
+    got = pep508_env(os = "exotic", arch = "exotic", python_version = "3.1.1")
+    got.pop("_aliases")
+    env.expect.that_dict(got).contains_exactly({
+        "implementation_name": "cpython",
+        "implementation_version": "3.1.1",
+        "os_name": "posix",
+        "platform_machine": "",
+        "platform_python_implementation": "CPython",
+        "platform_release": "",
+        "platform_system": "",
+        "platform_version": "0",
+        "python_full_version": "3.1.1",
+        "python_version": "3.1",
+        "sys_platform": "",
+    })
+
+_tests.append(_test_env_defaults)
+
+def _test_env_freebsd(env):
+    got = pep508_env(os = "freebsd", arch = "arm64", python_version = "3.1.1")
+    got.pop("_aliases")
+    env.expect.that_dict(got).contains_exactly({
+        "implementation_name": "cpython",
+        "implementation_version": "3.1.1",
+        "os_name": "posix",
+        "platform_machine": "aarch64",
+        "platform_python_implementation": "CPython",
+        "platform_release": "",
+        "platform_system": "FreeBSD",
+        "platform_version": "0",
+        "python_full_version": "3.1.1",
+        "python_version": "3.1",
+        "sys_platform": "freebsd",
+    })
+
+_tests.append(_test_env_freebsd)
+
+def _test_env_macos(env):
+    got = pep508_env(os = "macos", arch = "arm64", python_version = "3.1.1")
+    got.pop("_aliases")
+    env.expect.that_dict(got).contains_exactly({
+        "implementation_name": "cpython",
+        "implementation_version": "3.1.1",
+        "os_name": "posix",
+        "platform_machine": "aarch64",
+        "platform_python_implementation": "CPython",
+        "platform_release": "",
+        "platform_system": "Darwin",
+        "platform_version": "0",
+        "python_full_version": "3.1.1",
+        "python_version": "3.1",
+        "sys_platform": "darwin",
+    })
+
+_tests.append(_test_env_macos)
+
+def env_test_suite(name):  # buildifier: disable=function-docstring
+    test_suite(
+        name = name,
+        basic_tests = _tests,
+    )
diff --git a/tests/pypi/pep508/evaluate_tests.bzl b/tests/pypi/pep508/evaluate_tests.bzl
@@ -16,7 +16,6 @@
 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 = []
 
@@ -244,26 +243,37 @@ _tests.append(_evaluate_partial_only_extra)
 
 def _evaluate_with_aliases(env):
     # When
-    for target_platform, tests in {
+    for (os, cpu), tests in {
         # buildifier: @unsorted-dict-items
-        "osx_aarch64": {
+        ("osx", "aarch64"): {
             "platform_system == 'Darwin' and platform_machine == 'arm64'": True,
             "platform_system == 'Darwin' and platform_machine == 'aarch64'": True,
             "platform_system == 'Darwin' and platform_machine == 'amd64'": False,
         },
-        "osx_x86_64": {
+        ("osx", "x86_64"): {
             "platform_system == 'Darwin' and platform_machine == 'amd64'": True,
             "platform_system == 'Darwin' and platform_machine == 'x86_64'": True,
         },
-        "osx_x86_32": {
+        ("osx", "x86_32"): {
             "platform_system == 'Darwin' and platform_machine == 'i386'": True,
             "platform_system == 'Darwin' and platform_machine == 'i686'": True,
             "platform_system == 'Darwin' and platform_machine == 'x86_32'": True,
             "platform_system == 'Darwin' and platform_machine == 'x86_64'": False,
         },
+        ("freebsd", "x86_32"): {
+            "platform_system == 'FreeBSD' and platform_machine == 'i386'": True,
+            "platform_system == 'FreeBSD' and platform_machine == 'i686'": True,
+            "platform_system == 'FreeBSD' and platform_machine == 'x86_32'": True,
+            "platform_system == 'FreeBSD' and platform_machine == 'x86_64'": False,
+            "platform_system == 'FreeBSD' and os_name == 'posix'": True,
+        },
     }.items():  # buildifier: @unsorted-dict-items
         for input, want in tests.items():
-            _check_evaluate(env, input, want, pep508_env(platform_from_str(target_platform, "")))
+            _check_evaluate(env, input, want, pep508_env(
+                os = os,
+                arch = cpu,
+                python_version = "3.2",
+            ))
 
 _tests.append(_evaluate_with_aliases)
 
