third-party/benchmark/setup.py

   1 import contextlib
   2 import os
   3 import platform
   4 import shutil
   5 import sysconfig
   6 from pathlib import Path
   7 from typing import Generator
   8
   9 import setuptools
  10 from setuptools.command import build_ext
  11
  12 PYTHON_INCLUDE_PATH_PLACEHOLDER = "<PYTHON_INCLUDE_PATH>"
  13
  14 IS_WINDOWS = platform.system() == "Windows"
  15 IS_MAC = platform.system() == "Darwin"
  16
  17
  18 @contextlib.contextmanager
  19 def 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
  39 class 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
  50 class 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
 115 setuptools.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 )