• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import contextlib
2import os
3import platform
4import shutil
5import sysconfig
6from pathlib import Path
7
8import setuptools
9from setuptools.command import build_ext
10
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):
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        build_ext.build_ext.run(self)
57
58    def bazel_build(self, ext: BazelExtension):
59        """Runs the bazel build to create the package."""
60        with temp_fill_include_path("WORKSPACE"):
61            temp_path = Path(self.build_temp)
62
63            bazel_argv = [
64                "bazel",
65                "build",
66                ext.bazel_target,
67                f"--symlink_prefix={temp_path / 'bazel-'}",
68                f"--compilation_mode={'dbg' if self.debug else 'opt'}",
69                # C++17 is required by nanobind
70                f"--cxxopt={'/std:c++17' if IS_WINDOWS else '-std=c++17'}",
71            ]
72
73            if IS_WINDOWS:
74                # Link with python*.lib.
75                for library_dir in self.library_dirs:
76                    bazel_argv.append("--linkopt=/LIBPATH:" + library_dir)
77            elif IS_MAC:
78                if platform.machine() == "x86_64":
79                    # C++17 needs macOS 10.14 at minimum
80                    bazel_argv.append("--macos_minimum_os=10.14")
81
82                    # cross-compilation for Mac ARM64 on GitHub Mac x86 runners.
83                    # ARCHFLAGS is set by cibuildwheel before macOS wheel builds.
84                    archflags = os.getenv("ARCHFLAGS", "")
85                    if "arm64" in archflags:
86                        bazel_argv.append("--cpu=darwin_arm64")
87                        bazel_argv.append("--macos_cpus=arm64")
88
89                elif platform.machine() == "arm64":
90                    bazel_argv.append("--macos_minimum_os=11.0")
91
92            self.spawn(bazel_argv)
93
94            shared_lib_suffix = '.dll' if IS_WINDOWS else '.so'
95            ext_name = ext.target_name + shared_lib_suffix
96            ext_bazel_bin_path = temp_path / 'bazel-bin' / ext.relpath / ext_name
97
98            ext_dest_path = Path(self.get_ext_fullpath(ext.name))
99            shutil.copyfile(ext_bazel_bin_path, ext_dest_path)
100
101            # explicitly call `bazel shutdown` for graceful exit
102            self.spawn(["bazel", "shutdown"])
103
104
105setuptools.setup(
106    cmdclass=dict(build_ext=BuildBazelExtension),
107    ext_modules=[
108        BazelExtension(
109            name="google_benchmark._benchmark",
110            bazel_target="//bindings/python/google_benchmark:_benchmark",
111        )
112    ],
113)
114