1"""Rules to create python distribution files and properly name them""" 2 3load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo") 4load("@system_python//:version.bzl", "SYSTEM_PYTHON_VERSION") 5 6def _get_suffix(limited_api, python_version, cpu): 7 """Computes an ABI version tag for an extension module per PEP 3149.""" 8 if "win32" in cpu or "win64" in cpu: 9 if limited_api: 10 return ".pyd" 11 if "win32" in cpu: 12 abi = "win32" 13 elif "win64" in cpu: 14 abi = "win_amd64" 15 else: 16 fail("Unsupported CPU: " + cpu) 17 return ".cp{}-{}.{}".format(python_version, abi, "pyd") 18 19 if python_version == "system": 20 python_version = SYSTEM_PYTHON_VERSION 21 if int(python_version) < 38: 22 python_version += "m" 23 abis = { 24 "darwin_arm64": "darwin", 25 "darwin_x86_64": "darwin", 26 "darwin": "darwin", 27 "osx-x86_64": "darwin", 28 "osx-aarch_64": "darwin", 29 "linux-aarch_64": "aarch64-linux-gnu", 30 "linux-x86_64": "x86_64-linux-gnu", 31 "k8": "x86_64-linux-gnu", 32 } 33 34 return ".cpython-{}-{}.{}".format( 35 python_version, 36 abis[cpu], 37 "so" if limited_api else "abi3.so", 38 ) 39 elif limited_api: 40 return ".abi3.so" 41 42 fail("Unsupported combination of flags") 43 44def _declare_module_file(ctx, module_name, python_version, limited_api): 45 """Declares an output file for a Python module with this name, version, and limited api.""" 46 base_filename = module_name.replace(".", "/") 47 suffix = _get_suffix( 48 python_version = python_version, 49 limited_api = limited_api, 50 cpu = ctx.var["TARGET_CPU"], 51 ) 52 filename = base_filename + suffix 53 return ctx.actions.declare_file(filename) 54 55# -------------------------------------------------------------------------------------------------- 56# py_dist_module() 57# 58# Creates a Python binary extension module that is ready for distribution. 59# 60# py_dist_module( 61# name = "message_mod", 62# extension = "//python:_message_binary", 63# module_name = "google._upb._message", 64# ) 65# 66# In the simple case, this simply involves copying the input file to the proper filename for 67# our current configuration (module_name, cpu, python_version, limited_abi). 68# 69# For multiarch platforms (osx-universal2), we must combine binaries for multiple architectures 70# into a single output binary using the "llvm-lipo" tool. A config transition depends on multiple 71# architectures to get us the input files we need. 72 73def _py_multiarch_transition_impl(settings, attr): 74 if settings["//command_line_option:cpu"] == "osx-universal2": 75 return [{"//command_line_option:cpu": cpu} for cpu in ["osx-aarch_64", "osx-x86_64"]] 76 else: 77 return settings 78 79_py_multiarch_transition = transition( 80 implementation = _py_multiarch_transition_impl, 81 inputs = ["//command_line_option:cpu"], 82 outputs = ["//command_line_option:cpu"], 83) 84 85def _py_dist_module_impl(ctx): 86 output_file = _declare_module_file( 87 ctx = ctx, 88 module_name = ctx.attr.module_name, 89 python_version = ctx.attr._python_version[BuildSettingInfo].value, 90 limited_api = ctx.attr._limited_api[BuildSettingInfo].value, 91 ) 92 if len(ctx.attr.extension) == 1: 93 src = ctx.attr.extension[0][DefaultInfo].files.to_list()[0] 94 ctx.actions.run( 95 executable = "cp", 96 arguments = [src.path, output_file.path], 97 inputs = [src], 98 outputs = [output_file], 99 ) 100 return [ 101 DefaultInfo(files = depset([output_file])), 102 ] 103 else: 104 srcs = [mod[DefaultInfo].files.to_list()[0] for mod in ctx.attr.extension] 105 ctx.actions.run( 106 executable = "/usr/local/bin/llvm-lipo", 107 arguments = ["-create", "-output", output_file.path] + [src.path for src in srcs], 108 inputs = srcs, 109 outputs = [output_file], 110 ) 111 return [ 112 DefaultInfo(files = depset([output_file])), 113 ] 114 115py_dist_module = rule( 116 implementation = _py_dist_module_impl, 117 attrs = { 118 "module_name": attr.string(mandatory = True), 119 "extension": attr.label( 120 mandatory = True, 121 cfg = _py_multiarch_transition, 122 ), 123 "_limited_api": attr.label(default = "//python:limited_api"), 124 "_python_version": attr.label(default = "//python:python_version"), 125 "_allowlist_function_transition": attr.label( 126 default = "@bazel_tools//tools/allowlists/function_transition_allowlist", 127 ), 128 }, 129) 130 131# -------------------------------------------------------------------------------------------------- 132# py_dist() 133# 134# A rule that builds a collection of binary wheels, using transitions to depend on many different 135# python versions and cpus. 136 137def _py_dist_transition_impl(settings, attr): 138 _ignore = (settings) # @unused 139 transitions = [] 140 141 for cpu, version in attr.limited_api_wheels.items(): 142 transitions.append({ 143 "//command_line_option:cpu": cpu, 144 "//python:python_version": version, 145 "//python:limited_api": True, 146 }) 147 148 for version in attr.full_api_versions: 149 for cpu in attr.full_api_cpus: 150 transitions.append({ 151 "//command_line_option:cpu": cpu, 152 "//python:python_version": version, 153 "//python:limited_api": False, 154 }) 155 156 return transitions 157 158_py_dist_transition = transition( 159 implementation = _py_dist_transition_impl, 160 inputs = [], 161 outputs = [ 162 "//command_line_option:cpu", 163 "//python:python_version", 164 "//python:limited_api", 165 ], 166) 167 168def _py_dist_impl(ctx): 169 binary_files = [dep[DefaultInfo].files for dep in ctx.attr.binary_wheel] 170 pure_python_files = [ctx.attr.pure_python_wheel[DefaultInfo].files] 171 return [ 172 DefaultInfo(files = depset( 173 transitive = binary_files + pure_python_files, 174 )), 175 ] 176 177py_dist = rule( 178 implementation = _py_dist_impl, 179 attrs = { 180 "binary_wheel": attr.label( 181 mandatory = True, 182 cfg = _py_dist_transition, 183 ), 184 "pure_python_wheel": attr.label(mandatory = True), 185 "limited_api_wheels": attr.string_dict(), 186 "full_api_versions": attr.string_list(), 187 "full_api_cpus": attr.string_list(), 188 "_allowlist_function_transition": attr.label( 189 default = "@bazel_tools//tools/allowlists/function_transition_allowlist", 190 ), 191 }, 192) 193