1import contextlib 2import os 3import platform 4import shutil 5import sysconfig 6from pathlib import Path 7from typing import List 8 9import setuptools 10from setuptools.command import build_ext 11 12 13PYTHON_INCLUDE_PATH_PLACEHOLDER = "<PYTHON_INCLUDE_PATH>" 14 15IS_WINDOWS = platform.system() == "Windows" 16IS_MAC = platform.system() == "Darwin" 17 18 19def _get_long_description(fp: str) -> str: 20 with open(fp, "r", encoding="utf-8") as f: 21 return f.read() 22 23 24def _get_version(fp: str) -> str: 25 """Parse a version string from a file.""" 26 with open(fp, "r") as f: 27 for line in f: 28 if "__version__" in line: 29 delim = '"' 30 return line.split(delim)[1] 31 raise RuntimeError(f"could not find a version string in file {fp!r}.") 32 33 34def _parse_requirements(fp: str) -> List[str]: 35 with open(fp) as requirements: 36 return [ 37 line.rstrip() 38 for line in requirements 39 if not (line.isspace() or line.startswith("#")) 40 ] 41 42 43@contextlib.contextmanager 44def temp_fill_include_path(fp: str): 45 """Temporarily set the Python include path in a file.""" 46 with open(fp, "r+") as f: 47 try: 48 content = f.read() 49 replaced = content.replace( 50 PYTHON_INCLUDE_PATH_PLACEHOLDER, 51 Path(sysconfig.get_paths()['include']).as_posix(), 52 ) 53 f.seek(0) 54 f.write(replaced) 55 f.truncate() 56 yield 57 finally: 58 # revert to the original content after exit 59 f.seek(0) 60 f.write(content) 61 f.truncate() 62 63 64class BazelExtension(setuptools.Extension): 65 """A C/C++ extension that is defined as a Bazel BUILD target.""" 66 67 def __init__(self, name: str, bazel_target: str): 68 super().__init__(name=name, sources=[]) 69 70 self.bazel_target = bazel_target 71 stripped_target = bazel_target.split("//")[-1] 72 self.relpath, self.target_name = stripped_target.split(":") 73 74 75class BuildBazelExtension(build_ext.build_ext): 76 """A command that runs Bazel to build a C/C++ extension.""" 77 78 def run(self): 79 for ext in self.extensions: 80 self.bazel_build(ext) 81 build_ext.build_ext.run(self) 82 83 def bazel_build(self, ext: BazelExtension): 84 """Runs the bazel build to create the package.""" 85 with temp_fill_include_path("WORKSPACE"): 86 temp_path = Path(self.build_temp) 87 88 bazel_argv = [ 89 "bazel", 90 "build", 91 ext.bazel_target, 92 f"--symlink_prefix={temp_path / 'bazel-'}", 93 f"--compilation_mode={'dbg' if self.debug else 'opt'}", 94 # C++17 is required by nanobind 95 f"--cxxopt={'/std:c++17' if IS_WINDOWS else '-std=c++17'}", 96 ] 97 98 if IS_WINDOWS: 99 # Link with python*.lib. 100 for library_dir in self.library_dirs: 101 bazel_argv.append("--linkopt=/LIBPATH:" + library_dir) 102 elif IS_MAC: 103 if platform.machine() == "x86_64": 104 # C++17 needs macOS 10.14 at minimum 105 bazel_argv.append("--macos_minimum_os=10.14") 106 107 # cross-compilation for Mac ARM64 on GitHub Mac x86 runners. 108 # ARCHFLAGS is set by cibuildwheel before macOS wheel builds. 109 archflags = os.getenv("ARCHFLAGS", "") 110 if "arm64" in archflags: 111 bazel_argv.append("--cpu=darwin_arm64") 112 bazel_argv.append("--macos_cpus=arm64") 113 114 elif platform.machine() == "arm64": 115 bazel_argv.append("--macos_minimum_os=11.0") 116 117 self.spawn(bazel_argv) 118 119 shared_lib_suffix = '.dll' if IS_WINDOWS else '.so' 120 ext_name = ext.target_name + shared_lib_suffix 121 ext_bazel_bin_path = temp_path / 'bazel-bin' / ext.relpath / ext_name 122 123 ext_dest_path = Path(self.get_ext_fullpath(ext.name)) 124 shutil.copyfile(ext_bazel_bin_path, ext_dest_path) 125 126 # explicitly call `bazel shutdown` for graceful exit 127 self.spawn(["bazel", "shutdown"]) 128 129 130setuptools.setup( 131 name="google_benchmark", 132 version=_get_version("bindings/python/google_benchmark/__init__.py"), 133 url="https://github.com/google/benchmark", 134 description="A library to benchmark code snippets.", 135 long_description=_get_long_description("README.md"), 136 long_description_content_type="text/markdown", 137 author="Google", 138 author_email="benchmark-py@google.com", 139 # Contained modules and scripts. 140 package_dir={"": "bindings/python"}, 141 packages=setuptools.find_packages("bindings/python"), 142 install_requires=_parse_requirements("bindings/python/requirements.txt"), 143 cmdclass=dict(build_ext=BuildBazelExtension), 144 ext_modules=[ 145 BazelExtension( 146 "google_benchmark._benchmark", 147 "//bindings/python/google_benchmark:_benchmark", 148 ) 149 ], 150 zip_safe=False, 151 # PyPI package information. 152 classifiers=[ 153 "Development Status :: 4 - Beta", 154 "Intended Audience :: Developers", 155 "Intended Audience :: Science/Research", 156 "License :: OSI Approved :: Apache Software License", 157 "Programming Language :: Python :: 3.8", 158 "Programming Language :: Python :: 3.9", 159 "Programming Language :: Python :: 3.10", 160 "Programming Language :: Python :: 3.11", 161 "Topic :: Software Development :: Testing", 162 "Topic :: System :: Benchmark", 163 ], 164 license="Apache 2.0", 165 keywords="benchmark", 166) 167