Fixed installation from source.

grain/_src/python/experimental/index_shuffle/python/BUILD

@@ -7,10 +7,7 @@ licenses(["notice"])
 pybind_extension(
     name = "index_shuffle_module",
     srcs = ["index_shuffle_module.cc"],
-    deps = [
-        "//grain/_src/python/experimental/index_shuffle",
-        "@abseil-cpp//absl/strings",
-    ],
+    deps = ["//grain/_src/python/experimental/index_shuffle"],
 )
 
 py_test(

grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc

@@ -1,8 +1,8 @@
 #include <pybind11/pybind11.h>
 
 #include <cstdint>
+#include <string>
 
-#include "absl/strings/str_cat.h"
 #include "grain/_src/python/experimental/index_shuffle/index_shuffle.h"
 
 namespace py = pybind11;
@@ -16,15 +16,18 @@ PYBIND11_MODULE(index_shuffle_module, m) {
       "index_shuffle",
       [](int64_t index, int64_t max_index, uint32_t seed, uint32_t rounds) {
         if (rounds < 4 || rounds % 2 != 0 || rounds > 1024) {
+          // Using std::to_string to avoid extra dependency in setup.py.
           throw py::value_error(
-              absl::StrCat("rounds must be an even integer between 4 and 1024, "
-                           "but got rounds = ",
-                           rounds));
+              "rounds must be an even integer between 4 and 1024, "
+              "but got rounds = " +
+              std::to_string(rounds));
         }
         if (index < 0 || index > max_index) {
-          throw py::value_error(absl::StrCat(
-              "index must be in [0, max_index], but got index = ", index,
-              " and max_index = ", max_index));
+          // Using std::to_string to avoid extra dependency in setup.py.
+          throw py::value_error(
+              "index must be in [0, max_index], but got index = " +
+              std::to_string(index) +
+              " and max_index = " + std::to_string(max_index));
         }
         return grain::random::index_shuffle(index, max_index, seed, rounds);
       },

grain/oss/build_whl.sh

@@ -112,7 +112,7 @@ main() {
   if [ "$(uname)" == "Darwin" ]; then
     WHEEL_BLD_ARGS="${WHEEL_BLD_ARGS} --plat-name macosx_11_0_$(uname -m)"
   fi
-  "$PYTHON_BIN" setup.py $WHEEL_BLD_ARGS
+  GRAIN_SKIP_EXTRA_BUILD=1 "$PYTHON_BIN" setup.py $WHEEL_BLD_ARGS
 
   if [ -n "${AUDITWHEEL_PLATFORM}" ]; then
     printf '%s : "=== Auditing wheel\n' "$(date)"

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools"]
+requires = ["setuptools", "grpcio-tools", "pybind11"]
 build-backend = "setuptools.build_meta"
 
 [project]

setup.py

@@ -2,10 +2,19 @@
 
 Most project configs are in `pyproject.toml` -- prefer to modify
 `pyproject.toml` over this file if possible.
+
+There are extra building steps in this script:
+  1. Generates proto bindings.
+  2. Compiles C++ extensions using pybind11.
+In case those are pre-built (e.g. when building wheels), you can skip this
+extra build using `GRAIN_SKIP_EXTRA_BUILD=1`.
 """
 
+import importlib.resources
+import os
 import setuptools
 from setuptools import dist
+from setuptools.command import build_py
 
 
 class BinaryDistribution(dist.Distribution):
@@ -15,6 +24,78 @@ def has_ext_modules(self):
     return True
 
 
-setuptools.setup(
-    distclass=BinaryDistribution,
-)
+# If GRAIN_SKIP_EXTRA_BUILD=1, we are packaging pre-built binaries (e.g. in
+# build_whl.sh). We can skip all build logic (C++ compilation and Proto
+# generation).
+if os.environ.get("GRAIN_SKIP_EXTRA_BUILD", "0") == "1":
+  setuptools.setup(
+      distclass=BinaryDistribution,
+  )
+else:
+  # Full build logic for installing from source
+
+  class GenerateProtosCommand(setuptools.Command):
+    """Command to generate Python protobuf bindings."""
+
+    description = "Generate Python protobuf bindings"
+    user_options = []
+
+    def initialize_options(self):
+      pass
+
+    def finalize_options(self):
+      pass
+
+    def run(self):
+      from grpc_tools import protoc  # pylint: disable=g-import-not-at-top
+
+      root_dir = os.path.dirname(os.path.abspath(__file__))
+
+      proto_file = os.path.join(
+          root_dir, "grain", "proto", "execution_summary.proto"
+      )
+
+      grpc_protos_include = str(
+          importlib.resources.files("grpc_tools").joinpath("_proto")
+      )
+
+      proto_args = [
+          "grpc_tools.protoc",
+          f"--proto_path={grpc_protos_include}",
+          f"--proto_path={root_dir}",
+          f"--python_out={root_dir}",
+          proto_file,
+      ]
+
+      if protoc.main(proto_args) != 0:
+        raise RuntimeError(f"Error compiling proto: {proto_args}")
+
+  class BuildPyCommand(build_py.build_py):
+    """Run proto generation before build_py."""
+
+    def run(self):
+      self.run_command("generate_protos")
+      super().run()
+
+  from pybind11.setup_helpers import Pybind11Extension  # pylint: disable=g-import-not-at-top
+
+  ext_modules = [
+      Pybind11Extension(
+          name="grain._src.python.experimental.index_shuffle.python.index_shuffle_module",
+          sources=[
+              "grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc",
+              "grain/_src/python/experimental/index_shuffle/index_shuffle.cc",
+          ],
+          include_dirs=["."],
+      )
+  ]
+
+  setuptools.setup(
+      distclass=BinaryDistribution,
+      cmdclass={
+          "build_py": BuildPyCommand,
+          "generate_protos": GenerateProtosCommand,
+      },
+      ext_modules=ext_modules,
+      setup_requires=["grpcio-tools"],
+  )