From 4a2342be20a5fc679b9ba448296568748f69c694 Mon Sep 17 00:00:00 2001
From: Nicholas Junge <nicholas.junge@web.de>
Date: Mon, 29 Jul 2024 23:16:32 +0200
Subject: [PATCH] Bump nanobind-bazel to v2.1.0, add stubgen target

Adds a stub file for the `google_benchmark._benchmark` submodule,
generated with the new `nanobind_stubgen` rule released in nanobind_bazel
v2.1.0.

Tweaks the setup.py logic a little bit to package stub files with the
rest of the build artifacts. Also explicitly adds the generated stub and
marker files to the list of package data artifacts.
---
 MODULE.bazel                           |  2 +-
 bindings/python/google_benchmark/BUILD |  8 +++++-
 setup.py                               | 38 ++++++++++++++++++++------
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/MODULE.bazel b/MODULE.bazel
index e368e325c7..4210ea0be2 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -38,4 +38,4 @@ use_repo(pip, "tools_pip_deps")
 
 # -- bazel_dep definitions -- #
 
-bazel_dep(name = "nanobind_bazel", version = "2.0.0", dev_dependency = True)
+bazel_dep(name = "nanobind_bazel", version = "2.1.0", dev_dependency = True)
diff --git a/bindings/python/google_benchmark/BUILD b/bindings/python/google_benchmark/BUILD
index 0c8e3c103f..30e389337d 100644
--- a/bindings/python/google_benchmark/BUILD
+++ b/bindings/python/google_benchmark/BUILD
@@ -1,4 +1,4 @@
-load("@nanobind_bazel//:build_defs.bzl", "nanobind_extension")
+load("@nanobind_bazel//:build_defs.bzl", "nanobind_extension", "nanobind_stubgen")
 
 py_library(
     name = "google_benchmark",
@@ -15,6 +15,12 @@ nanobind_extension(
     deps = ["//:benchmark"],
 )
 
+nanobind_stubgen(
+    name = "benchmark_stubgen",
+    marker_file = "bindings/python/google_benchmark/py.typed",
+    module = ":_benchmark",
+)
+
 py_test(
     name = "example",
     srcs = ["example.py"],
diff --git a/setup.py b/setup.py
index 40cdc8d339..d171476f7e 100644
--- a/setup.py
+++ b/setup.py
@@ -99,7 +99,7 @@ def bazel_build(self, ext: BazelExtension) -> None:
 
         bazel_argv = [
             "bazel",
-            "build",
+            "run",
             ext.bazel_target,
             f"--symlink_prefix={temp_path / 'bazel-'}",
             f"--compilation_mode={'dbg' if self.debug else 'opt'}",
@@ -127,20 +127,42 @@ def bazel_build(self, ext: BazelExtension) -> None:
         else:
             suffix = ".abi3.so" if ext.py_limited_api else ".so"
 
-        ext_name = ext.target_name + suffix
-        ext_bazel_bin_path = temp_path / "bazel-bin" / ext.relpath / ext_name
-        ext_dest_path = Path(self.get_ext_fullpath(ext.name)).with_name(
-            ext_name
-        )
-        shutil.copyfile(ext_bazel_bin_path, ext_dest_path)
+        # copy the Bazel build artifacts into setuptools' libdir,
+        # from where the wheel is built.
+        pkgname = "google_benchmark"
+        pythonroot = Path("bindings") / "python" / "google_benchmark"
+        srcdir = temp_path / "bazel-bin" / pythonroot
+        libdir = Path(self.build_lib) / pkgname
+        for root, dirs, files in os.walk(srcdir, topdown=True):
+            # exclude runfiles directories and children.
+            dirs[:] = [d for d in dirs if "runfiles" not in d]
+
+            for f in files:
+                print(f)
+                fp = Path(f)
+                should_copy = False
+                # we do not want the bare .so file included
+                # when building for ABI3, so we require a
+                # full and exact match on the file extension.
+                if "".join(fp.suffixes) == suffix:
+                    should_copy = True
+                elif fp.suffix == ".pyi":
+                    should_copy = True
+                elif Path(root) == srcdir and f == "py.typed":
+                    # copy py.typed, but only at the package root.
+                    should_copy = True
+
+                if should_copy:
+                    shutil.copyfile(root / fp, libdir / fp)
 
 
 setuptools.setup(
     cmdclass=dict(build_ext=BuildBazelExtension),
+    package_data={"google_benchmark": ["py.typed", "*.pyi"]},
     ext_modules=[
         BazelExtension(
             name="google_benchmark._benchmark",
-            bazel_target="//bindings/python/google_benchmark:_benchmark",
+            bazel_target="//bindings/python/google_benchmark:benchmark_stubgen",
             py_limited_api=py_limited_api,
         )
     ],