|
| 1 | +# Copyright 2026 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Helper for building CEL Python extensions with bazel for PyPI releases.""" |
| 16 | + |
| 17 | +import glob |
| 18 | +import os |
| 19 | +import platform |
| 20 | +import re |
| 21 | +import shutil |
| 22 | +import subprocess |
| 23 | +import sys |
| 24 | +import setuptools |
| 25 | +import setuptools.command.build_ext |
| 26 | + |
| 27 | + |
| 28 | +class BazelExtension(setuptools.Extension): |
| 29 | + |
| 30 | + def __init__(self, name, target): |
| 31 | + super().__init__(name, sources=[]) |
| 32 | + self.target = target |
| 33 | + |
| 34 | + |
| 35 | +class BazelBuild(setuptools.command.build_ext.build_ext): |
| 36 | + """Custom build_ext command to build extensions with bazel.""" |
| 37 | + |
| 38 | + def run(self): |
| 39 | + # Ensure bazel is available |
| 40 | + if shutil.which('bazel') is None: |
| 41 | + raise RuntimeError('bazel must be installed to build the extensions.') |
| 42 | + |
| 43 | + for ext in self.extensions: |
| 44 | + self.build_extension(ext) |
| 45 | + |
| 46 | + def build_extension(self, ext): |
| 47 | + # cibuildwheel executes setup.py using the target Python version. |
| 48 | + # Thus, we can use the current Python version as the target version for |
| 49 | + # the bazel build. |
| 50 | + python_version = f'{sys.version_info.major}.{sys.version_info.minor}' |
| 51 | + |
| 52 | + print(f'Building for target Python version: {python_version}') |
| 53 | + |
| 54 | + module_bazel_path = os.path.join( |
| 55 | + os.path.dirname(os.path.abspath(__file__)), 'MODULE.bazel' |
| 56 | + ) |
| 57 | + if os.path.exists(module_bazel_path): |
| 58 | + with open(module_bazel_path, 'r') as f: |
| 59 | + content = f.read() |
| 60 | + content = re.sub( |
| 61 | + r'python_version\s*=\s*".*"', |
| 62 | + f'python_version = "{python_version}"', |
| 63 | + content, |
| 64 | + ) |
| 65 | + with open(module_bazel_path, 'w') as f: |
| 66 | + f.write(content) |
| 67 | + else: |
| 68 | + raise RuntimeError(f'MODULE.bazel not found at {module_bazel_path}') |
| 69 | + |
| 70 | + dest_path = self.get_ext_fullpath(ext.name) |
| 71 | + dest_dir = os.path.dirname(dest_path) |
| 72 | + os.makedirs(dest_dir, exist_ok=True) |
| 73 | + |
| 74 | + # Build with bazel |
| 75 | + # Use --compilation_mode=opt for release builds |
| 76 | + cmd = [ |
| 77 | + 'bazel', |
| 78 | + 'build', |
| 79 | + ext.target, |
| 80 | + '--compilation_mode=opt', |
| 81 | + '--symlink_prefix=bazel-', |
| 82 | + ] |
| 83 | + extra_flags = os.environ.get('CEL_BAZEL_FLAGS') |
| 84 | + if extra_flags: |
| 85 | + cmd.extend(extra_flags.split()) |
| 86 | + if sys.platform == 'win32': |
| 87 | + self.platform_config_windows(cmd, python_version) |
| 88 | + if sys.platform == 'darwin': |
| 89 | + self.platform_config_macos(cmd) |
| 90 | + print(f"Building {ext.name} with bazel: {' '.join(cmd)}") |
| 91 | + subprocess.check_call(cmd) |
| 92 | + |
| 93 | + # Determine the output path of the bazel build |
| 94 | + # We handle targets like //:py_cel and //ext:ext_math |
| 95 | + rel_path = ext.target.lstrip('/').replace(':', '/') |
| 96 | + if rel_path.startswith('/'): |
| 97 | + rel_path = rel_path[1:] |
| 98 | + |
| 99 | + # Bazel output directory |
| 100 | + bazel_bin = 'bazel-bin' |
| 101 | + if not os.path.exists(bazel_bin): |
| 102 | + try: |
| 103 | + bazel_bin = subprocess.check_output( |
| 104 | + ['bazel', 'info', 'bazel-bin', '-c', 'opt'], text=True |
| 105 | + ).strip() |
| 106 | + except Exception: # pylint: disable=broad-exception-caught |
| 107 | + pass |
| 108 | + |
| 109 | + # Attempt to find the generated shared library |
| 110 | + # Suffixes can vary by platform |
| 111 | + suffixes = ['.so', '.pyd', '.dylib'] |
| 112 | + found = None |
| 113 | + |
| 114 | + candidate_base = os.path.join(bazel_bin, rel_path) |
| 115 | + |
| 116 | + for suffix in suffixes: |
| 117 | + candidate = candidate_base + suffix |
| 118 | + if os.path.exists(candidate): |
| 119 | + found = candidate |
| 120 | + break |
| 121 | + |
| 122 | + if not found: |
| 123 | + # Fallback check for exact name match |
| 124 | + # (less likely for extensions but possible) |
| 125 | + if os.path.exists(candidate_base): |
| 126 | + found = candidate_base |
| 127 | + |
| 128 | + if not found: |
| 129 | + raise RuntimeError( |
| 130 | + f'Could not find bazel output for {ext.target} at {candidate_base}.*' |
| 131 | + ) |
| 132 | + |
| 133 | + print(f'Copying {found} to {dest_path}') |
| 134 | + shutil.copyfile(found, dest_path) |
| 135 | + |
| 136 | + def platform_config_windows(self, cmd, python_version): |
| 137 | + """Applies Windows-specific Bazel workarounds for Hermetic Python.""" |
| 138 | + cmd.insert(1, '--output_user_root=C:/tmp') |
| 139 | + |
| 140 | + # 1. Get output base |
| 141 | + output_base = subprocess.check_output( |
| 142 | + ['bazel', '--output_user_root=C:/tmp', 'info', 'output_base'], text=True |
| 143 | + ).strip() |
| 144 | + |
| 145 | + # 2. Find hermetic python directory |
| 146 | + py_ver_underscore = python_version.replace('.', '_') |
| 147 | + pattern = os.path.join( |
| 148 | + output_base, 'external', f'*python_{py_ver_underscore}_host' |
| 149 | + ) |
| 150 | + matches = glob.glob(pattern) |
| 151 | + if not matches: |
| 152 | + print( |
| 153 | + 'Warning: Hermetic Python directory not found with pattern:' |
| 154 | + f' {pattern}' |
| 155 | + ) |
| 156 | + return |
| 157 | + |
| 158 | + py_host_dir = matches[0] |
| 159 | + print(f'Found Hermetic Python Directory: {py_host_dir}') |
| 160 | + |
| 161 | + # 3. Copy libs to a space-free directory |
| 162 | + target_dir = r'C:\tmp\python_libs' |
| 163 | + os.makedirs(target_dir, exist_ok=True) |
| 164 | + |
| 165 | + lib_pattern = os.path.join(py_host_dir, 'libs', 'python*.lib') |
| 166 | + for lib_file in glob.glob(lib_pattern): |
| 167 | + print(f'Copying {lib_file} to {target_dir}') |
| 168 | + shutil.copy(lib_file, target_dir) |
| 169 | + |
| 170 | + dll_pattern = os.path.join(py_host_dir, 'python*.dll') |
| 171 | + for dll_file in glob.glob(dll_pattern): |
| 172 | + print(f'Copying {dll_file} to {target_dir}') |
| 173 | + shutil.copy(dll_file, target_dir) |
| 174 | + |
| 175 | + # 4. Add link flags to bazel command |
| 176 | + cmd.extend(['--linkopt=/LIBPATH:C:\\tmp\\python_libs', '--action_env=PATH']) |
| 177 | + |
| 178 | + # 5. Add to PATH in current process so bazel can find the DLLs |
| 179 | + os.environ['PATH'] = f"{target_dir};{os.environ.get('PATH', '')}" |
| 180 | + |
| 181 | + def platform_config_macos(self, cmd): |
| 182 | + """Applies macOS-specific Bazel configurations.""" |
| 183 | + deployment_target = os.environ.get('MACOSX_DEPLOYMENT_TARGET', '10.13') |
| 184 | + cmd.extend([ |
| 185 | + f'--macos_minimum_os={deployment_target}', |
| 186 | + '--cxxopt=-faligned-allocation', |
| 187 | + ]) |
| 188 | + |
| 189 | + archflags = os.environ.get('ARCHFLAGS', '') |
| 190 | + if 'x86_64' in archflags: |
| 191 | + target_arch = 'x86_64' |
| 192 | + elif 'arm64' in archflags: |
| 193 | + target_arch = 'arm64' |
| 194 | + else: |
| 195 | + machine = platform.machine() |
| 196 | + if machine in ('AMD64', 'x86_64'): |
| 197 | + target_arch = 'x86_64' |
| 198 | + elif machine in ('arm64', 'aarch64'): |
| 199 | + target_arch = 'arm64' |
| 200 | + else: |
| 201 | + target_arch = machine |
| 202 | + |
| 203 | + print(f'Target architecture for macOS: {target_arch}', flush=True) |
| 204 | + cmd.append(f'--macos_cpus={target_arch}') |
| 205 | + cmd.append(f'--cpu=darwin_{target_arch}') |
| 206 | + if target_arch == 'x86_64': |
| 207 | + cmd.append('--platforms=//cel_expr_python:macos_x86_64') |
0 commit comments