For macOS, build a single multiarch "universal2" wheel instead of separate x86_64 and arm64 wheels.

When I tried the previous `arm64` wheel on macOS, I discovered that `pip` on macOS only supports the `arm64` platform tag in a limited set of circumstances.  pip seems to prefer `universal2` wheels.

To build a `universal2` wheel, we must run the `llvm-lipo` tool to bundle multiple `cc_binary()` outputs into a single output wheel.  We use a transition to depend on multiple architectures for the extension, if we see that we are building for a multiarch CPU.

PiperOrigin-RevId: 447486256

python/dist/BUILD.bazel

@@ -76,6 +76,11 @@ config_setting(
     values = {"cpu": "osx-aarch_64"},
 )
 
+config_setting(
+    name = "osx-universal2_cpu",
+    values = {"cpu": "osx-universal2"},
+)
+
 config_setting(
     name = "win32_cpu",
     values = {"cpu": "win32"},
@@ -98,8 +103,7 @@ py_wheel(
     platform = select({
         ":x86_64_cpu": "manylinux2014_x86_64",
         ":aarch64_cpu": "manylinux2014_aarch64",
-        ":osx-x86_64_cpu": "macosx_10_9_x86_64",
+        ":osx-universal2_cpu": "macosx_10_9_universal2",
-        ":osx-aarch64_cpu": "macosx_10_9_arm64",
         ":win32_cpu": "win32",
         ":win64_cpu": "win_amd64",
         "//conditions:default": "any",
@@ -182,8 +186,7 @@ py_dist(
         "win64": "310",
         "linux-x86_64": "37",
         "linux-aarch_64": "37",
-        "osx-aarch_64": "37",
+        "osx-universal2": "37",
-        "osx-x86_64": "37",
     },
     pure_python_wheel = ":pure_python_wheel",
     tags = ["manual"],

python/dist/dist.bzl

@@ -4,6 +4,7 @@ load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
 load("@system_python//:version.bzl", "SYSTEM_PYTHON_VERSION")
 
 def _get_suffix(limited_api, python_version, cpu):
+    """Computes an ABI version tag for an extension module per PEP 3149."""
     suffix = "pyd" if ("win" in cpu) else "so"
 
     if limited_api == True:
@@ -37,50 +38,99 @@ def _get_suffix(limited_api, python_version, cpu):
 
     fail("Unsupported combination of flags")
 
-def _py_dist_module_impl(ctx):
+def _declare_module_file(ctx, module_name, python_version, limited_api):
-    base_filename = ctx.attr.module_name.replace(".", "/")
+    """Declares an output file for a Python module with this name, version, and limited api."""
+    base_filename = module_name.replace(".", "/")
     suffix = _get_suffix(
-        limited_api = ctx.attr._limited_api[BuildSettingInfo].value,
+        python_version = python_version,
-        python_version = ctx.attr._python_version[BuildSettingInfo].value,
+        limited_api = limited_api,
         cpu = ctx.var["TARGET_CPU"],
     )
     filename = base_filename + suffix
-    file = ctx.actions.declare_file(filename)
+    return ctx.actions.declare_file(filename)
-    src = ctx.attr.extension[DefaultInfo].files.to_list()[0]
+
-    ctx.actions.run(
+# --------------------------------------------------------------------------------------------------
-        executable = "cp",
+# py_dist_module()
-        arguments = [src.path, file.path],
+#
-        inputs = [src],
+# Creates a Python binary extension module that is ready for distribution.
-        outputs = [file],
+#
-    )
+#   py_dist_module(
-    return [
+#       name = "message_mod",
-        DefaultInfo(files = depset([file])),
+#       extension = "//python:_message_binary",
-    ]
+#       module_name = "google._upb._message",
+#   )
+#
+# In the simple case, this simply involves copying the input file to the proper filename for
+# our current configuration (module_name, cpu, python_version, limited_abi).
+#
+# For multiarch platforms (osx-universal2), we must combine binaries for multiple architectures
+# into a single output binary using the "llvm-lipo" tool.  A config transition depends on multiple
+# architectures to get us the input files we need.
+
+def _py_multiarch_transition_impl(settings, attr):
+    if settings["//command_line_option:cpu"] == "osx-universal2":
+        return [{"//command_line_option:cpu": cpu} for cpu in ["osx-aarch_64", "osx-x86_64"]]
+    else:
+        return settings
+
+_py_multiarch_transition = transition(
+    implementation = _py_multiarch_transition_impl,
+    inputs = ["//command_line_option:cpu"],
+    outputs = ["//command_line_option:cpu"],
+)
 
-_py_dist_module_rule = rule(
+def _py_dist_module_impl(ctx):
+    output_file = _declare_module_file(
+        ctx = ctx,
+        module_name = ctx.attr.module_name,
+        python_version = ctx.attr._python_version[BuildSettingInfo].value,
+        limited_api = ctx.attr._limited_api[BuildSettingInfo].value,
+    )
+    if len(ctx.attr.extension) == 1:
+        src = ctx.attr.extension[0][DefaultInfo].files.to_list()[0]
+        ctx.actions.run(
+            executable = "cp",
+            arguments = [src.path, output_file.path],
+            inputs = [src],
+            outputs = [output_file],
+        )
+        return [
+            DefaultInfo(files = depset([output_file])),
+        ]
+    else:
+        srcs = [mod[DefaultInfo].files.to_list()[0] for mod in ctx.attr.extension]
+        ctx.actions.run(
+            executable = "/usr/local/bin/llvm-lipo",
+            arguments = ["-create", "-output", output_file.path] + [src.path for src in srcs],
+            inputs = srcs,
+            outputs = [output_file],
+        )
+        return [
+            DefaultInfo(files = depset([output_file])),
+        ]
+
+py_dist_module = rule(
     output_to_genfiles = True,
     implementation = _py_dist_module_impl,
-    fragments = ["cpp"],
     attrs = {
         "module_name": attr.string(mandatory = True),
         "extension": attr.label(
             mandatory = True,
-            providers = [CcInfo],
+            cfg = _py_multiarch_transition,
         ),
         "_limited_api": attr.label(default = "//python:limited_api"),
         "_python_version": attr.label(default = "//python:python_version"),
-        "_cc_toolchain": attr.label(
+        "_allowlist_function_transition": attr.label(
-            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
+            default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
         ),
     },
 )
 
-def py_dist_module(name, module_name, extension):
+# --------------------------------------------------------------------------------------------------
-    _py_dist_module_rule(
+# py_dist()
-        name = name,
+#
-        module_name = module_name,
+# A rule that builds a collection of binary wheels, using transitions to depend on many different
-        extension = extension,
+# python versions and cpus.
-    )
 
 def _py_dist_transition_impl(settings, attr):
     _ignore = (settings)  # @unused