124 lines · python
1import contextlib2import os3import platform4import shutil5import sysconfig6from pathlib import Path7from typing import Generator8 9import setuptools10from setuptools.command import build_ext11 12PYTHON_INCLUDE_PATH_PLACEHOLDER = "<PYTHON_INCLUDE_PATH>"13 14IS_WINDOWS = platform.system() == "Windows"15IS_MAC = platform.system() == "Darwin"16 17 18@contextlib.contextmanager19def 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 yield32 finally:33 # revert to the original content after exit34 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_target46 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 exit58 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 to64 do again in the `build_ext` base class.65 """66 pass67 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 nanobind81 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 minimum91 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_suffix107 ext_bazel_bin_path = (108 temp_path / "bazel-bin" / ext.relpath / ext_name109 )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