9 files changed, 235 insertions, 0 deletions
diff --git a/bindings/python/BUILD b/bindings/python/BUILD
new file mode 100644
index 0000000..9559a76
--- /dev/null
+++ b/bindings/python/BUILD
@@ -0,0 +1,3 @@
+exports_files(glob(["*.BUILD"]))
+exports_files(["build_defs.bzl"])
+
diff --git a/bindings/python/benchmark/BUILD b/bindings/python/benchmark/BUILD
new file mode 100644
index 0000000..49f536e
--- /dev/null
+++ b/bindings/python/benchmark/BUILD
@@ -0,0 +1,38 @@
+load("//bindings/python:build_defs.bzl", "py_extension")
+
+py_library(
+    name = "benchmark",
+    srcs = ["__init__.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":_benchmark",
+        # pip; absl:app
+    ],
+)
+
+py_extension(
+    name = "_benchmark",
+    srcs = ["benchmark.cc"],
+    copts = [
+        "-fexceptions",
+        "-fno-strict-aliasing",
+    ],
+    features = ["-use_header_modules"],
+    deps = [
+        "//:benchmark",
+        "@pybind11",
+        "@python_headers",
+    ],
+)
+
+py_test(
+    name = "example",
+    srcs = ["example.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":benchmark",
+    ],
+)
+
diff --git a/bindings/python/benchmark/__init__.py b/bindings/python/benchmark/__init__.py
new file mode 100644
index 0000000..27f76e0
--- /dev/null
+++ b/bindings/python/benchmark/__init__.py
@@ -0,0 +1,62 @@
+# Copyright 2020 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Python benchmarking utilities.
+
+Example usage:
+  import benchmark
+
+  @benchmark.register
+  def my_benchmark(state):
+      ...  # Code executed outside `while` loop is not timed.
+
+      while state:
+        ...  # Code executed within `while` loop is timed.
+
+  if __name__ == '__main__':
+    benchmark.main()
+"""
+
+from absl import app
+from benchmark import _benchmark
+
+__all__ = [
+    "register",
+    "main",
+]
+
+__version__ = "0.1.0"
+
+
+def register(f=None, *, name=None):
+  if f is None:
+    return lambda f: register(f, name=name)
+  if name is None:
+    name = f.__name__
+  _benchmark.RegisterBenchmark(name, f)
+  return f
+
+
+def _flags_parser(argv):
+  argv = _benchmark.Initialize(argv)
+  return app.parse_flags_with_usage(argv)
+
+
+def _run_benchmarks(argv):
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+  return _benchmark.RunSpecifiedBenchmarks()
+
+
+def main(argv=None):
+  return app.run(_run_benchmarks, argv=argv, flags_parser=_flags_parser)
diff --git a/bindings/python/benchmark/benchmark.cc b/bindings/python/benchmark/benchmark.cc
new file mode 100644
index 0000000..ef95559
--- /dev/null
+++ b/bindings/python/benchmark/benchmark.cc
@@ -0,0 +1,47 @@
+// Benchmark for Python.
+
+#include "benchmark/benchmark.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/stl.h"
+
+namespace {
+namespace py = ::pybind11;
+
+std::vector<std::string> Initialize(const std::vector<std::string>& argv) {
+  // The `argv` pointers here become invalid when this function returns, but
+  // benchmark holds the pointer to `argv[0]`. We create a static copy of it
+  // so it persists, and replace the pointer below.
+  static std::string executable_name(argv[0]);
+  std::vector<char*> ptrs;
+  ptrs.reserve(argv.size());
+  for (auto& arg : argv) {
+    ptrs.push_back(const_cast<char*>(arg.c_str()));
+  }
+  ptrs[0] = const_cast<char*>(executable_name.c_str());
+  int argc = static_cast<int>(argv.size());
+  benchmark::Initialize(&argc, ptrs.data());
+  std::vector<std::string> remaining_argv;
+  remaining_argv.reserve(argc);
+  for (int i = 0; i < argc; ++i) {
+    remaining_argv.emplace_back(ptrs[i]);
+  }
+  return remaining_argv;
+}
+
+void RegisterBenchmark(const char* name, py::function f) {
+  benchmark::RegisterBenchmark(name, [f](benchmark::State& state) {
+    f(&state);
+  });
+}
+
+PYBIND11_MODULE(_benchmark, m) {
+  m.def("Initialize", Initialize);
+  m.def("RegisterBenchmark", RegisterBenchmark);
+  m.def("RunSpecifiedBenchmarks",
+        []() { benchmark::RunSpecifiedBenchmarks(); });
+
+  py::class_<benchmark::State>(m, "State")
+      .def("__bool__", &benchmark::State::KeepRunning)
+      .def_property_readonly("keep_running", &benchmark::State::KeepRunning);
+};
+}  // namespace
diff --git a/bindings/python/benchmark/example.py b/bindings/python/benchmark/example.py
new file mode 100644
index 0000000..24da127
--- /dev/null
+++ b/bindings/python/benchmark/example.py
@@ -0,0 +1,32 @@
+# Copyright 2020 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Example of Python using C++ benchmark framework."""
+
+import benchmark
+
+
+@benchmark.register
+def empty(state):
+  while state:
+    pass
+
+
+@benchmark.register
+def sum_million(state):
+  while state:
+    sum(range(1_000_000))
+
+
+if __name__ == '__main__':
+  benchmark.main()
diff --git a/bindings/python/build_defs.bzl b/bindings/python/build_defs.bzl
new file mode 100644
index 0000000..45907aa
--- /dev/null
+++ b/bindings/python/build_defs.bzl
@@ -0,0 +1,25 @@
+_SHARED_LIB_SUFFIX = {
+    "//conditions:default": ".so",
+    "//:windows": ".dll",
+}
+
+def py_extension(name, srcs, hdrs = [], copts = [], features = [], deps = []):
+    for shared_lib_suffix in _SHARED_LIB_SUFFIX.values():
+        shared_lib_name = name + shared_lib_suffix
+        native.cc_binary(
+            name = shared_lib_name,
+            linkshared = 1,
+            linkstatic = 1,
+            srcs = srcs + hdrs,
+            copts = copts,
+            features = features,
+            deps = deps,
+        )
+
+    return native.py_library(
+        name = name,
+        data = select({
+            platform: [name + shared_lib_suffix]
+            for platform, shared_lib_suffix in _SHARED_LIB_SUFFIX.items()
+        }),
+    )
diff --git a/bindings/python/pybind11.BUILD b/bindings/python/pybind11.BUILD
new file mode 100644
index 0000000..bc83350
--- /dev/null
+++ b/bindings/python/pybind11.BUILD
@@ -0,0 +1,20 @@
+cc_library(
+    name = "pybind11",
+    hdrs = glob(
+        include = [
+            "include/pybind11/*.h",
+            "include/pybind11/detail/*.h",
+        ],
+        exclude = [
+            "include/pybind11/common.h",
+            "include/pybind11/eigen.h",
+        ],
+    ),
+    copts = [
+        "-fexceptions",
+        "-Wno-undefined-inline",
+        "-Wno-pragma-once-outside-header",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+)
diff --git a/bindings/python/python_headers.BUILD b/bindings/python/python_headers.BUILD
new file mode 100644
index 0000000..9c34cf6
--- /dev/null
+++ b/bindings/python/python_headers.BUILD
@@ -0,0 +1,6 @@
+cc_library(
+    name = "python_headers",
+    hdrs = glob(["**/*.h"]),
+    includes = ["."],
+    visibility = ["//visibility:public"],
+)
diff --git a/bindings/python/requirements.txt b/bindings/python/requirements.txt
new file mode 100644
index 0000000..f5bbe7e
--- /dev/null
+++ b/bindings/python/requirements.txt
@@ -0,0 +1,2 @@
+absl-py>=0.7.1
+