1import contextlib 2import os 3import platform 4import shutil 5import sysconfig 6from pathlib import Path 7from typing import Generator 8 9import setuptools 10from setuptools.command import build_ext 11 12PYTHON_INCLUDE_PATH_PLACEHOLDER = "<PYTHON_INCLUDE_PATH>" 13 14IS_WINDOWS = platform.system() == "Windows" 15IS_MAC = platform.system() == "Darwin" 16 17 18@contextlib.contextmanager 19def temp_fill_include_path(fp: str) -> Generator[None, None, None]: 20 """Temporarily set the Python include path in a file.""" 21 with open(fp, "r+") as f: 22 try: 23 content = f.read() 24 replaced = content.replace( 25 PYTHON_INCLUDE_PATH_PLACEHOLDER, 26 Path(sysconfig.get_paths()["include"]).as_posix(), 27 ) 28 f.seek(0) 29 f.write(replaced) 30 f.truncate() 31 yield 32 finally: 33 # revert to the original content after exit 34 f.seek(0) 35 f.write(content) 36 f.truncate() 37 38 39class BazelExtension(setuptools.Extension): 40 """A C/C++ extension that is defined as a Bazel BUILD target.""" 41 42 def __init__(self, name: str, bazel_target: str): 43 super().__init__(name=name, sources=[]) 44 45 self.bazel_target = bazel_target 46 stripped_target = bazel_target.split("//")[-1] 47 self.relpath, self.target_name = stripped_target.split(":") 48 49 50class BuildBazelExtension(build_ext.build_ext): 51 """A command that runs Bazel to build a C/C++ extension.""" 52 53 def run(self): 54 for ext in self.extensions: 55 self.bazel_build(ext) 56 super().run() 57 # explicitly call `bazel shutdown` for graceful exit 58 self.spawn(["bazel", "shutdown"]) 59 60 def copy_extensions_to_source(self): 61 """ 62 Copy generated extensions into the source tree. 63 This is done in the ``bazel_build`` method, so it's not necessary to 64 do again in the `build_ext` base class. 65 """ 66 pass 67 68 def bazel_build(self, ext: BazelExtension) -> None: 69 """Runs the bazel build to create the package.""" 70 with temp_fill_include_path("WORKSPACE"): 71 temp_path = Path(self.build_temp) 72 73 bazel_argv = [ 74 "bazel", 75 "build", 76 ext.bazel_target, 77 "--enable_bzlmod=false", 78 f"--symlink_prefix={temp_path / 'bazel-'}", 79 f"--compilation_mode={'dbg' if self.debug else 'opt'}", 80 # C++17 is required by nanobind 81 f"--cxxopt={'/std:c++17' if IS_WINDOWS else '-std=c++17'}", 82 ] 83 84 if IS_WINDOWS: 85 # Link with python*.lib. 86 for library_dir in self.library_dirs: 87 bazel_argv.append("--linkopt=/LIBPATH:" + library_dir) 88 elif IS_MAC: 89 if platform.machine() == "x86_64": 90 # C++17 needs macOS 10.14 at minimum 91 bazel_argv.append("--macos_minimum_os=10.14") 92 93 # cross-compilation for Mac ARM64 on GitHub Mac x86 runners. 94 # ARCHFLAGS is set by cibuildwheel before macOS wheel builds. 95 archflags = os.getenv("ARCHFLAGS", "") 96 if "arm64" in archflags: 97 bazel_argv.append("--cpu=darwin_arm64") 98 bazel_argv.append("--macos_cpus=arm64") 99 100 elif platform.machine() == "arm64": 101 bazel_argv.append("--macos_minimum_os=11.0") 102 103 self.spawn(bazel_argv) 104 105 shared_lib_suffix = ".dll" if IS_WINDOWS else ".so" 106 ext_name = ext.target_name + shared_lib_suffix 107 ext_bazel_bin_path = ( 108 temp_path / "bazel-bin" / ext.relpath / ext_name 109 ) 110 111 ext_dest_path = Path(self.get_ext_fullpath(ext.name)) 112 shutil.copyfile(ext_bazel_bin_path, ext_dest_path) 113 114 115setuptools.setup( 116 cmdclass=dict(build_ext=BuildBazelExtension), 117 ext_modules=[ 118 BazelExtension( 119 name="google_benchmark._benchmark", 120 bazel_target="//bindings/python/google_benchmark:_benchmark", 121 ) 122 ], 123) 124