Add Python bindings. (#957)

* Add Python bindings.

* Add license headers.

* Change example to a test.

* Add example usage to module docstring.

WORKSPACE

@@ -13,3 +13,18 @@ http_archive(
     strip_prefix = "googletest-3f0cf6b62ad1eb50d8736538363d3580dd640c3e",
     urls = ["https://github.com/google/googletest/archive/3f0cf6b62ad1eb50d8736538363d3580dd640c3e.zip"],
 )
+
+http_archive(
+    name = "pybind11",
+    build_file = "@//bindings/python:pybind11.BUILD",
+    sha256 = "1eed57bc6863190e35637290f97a20c81cfe4d9090ac0a24f3bbf08f265eb71d",
+    strip_prefix = "pybind11-2.4.3",
+    urls = ["https://github.com/pybind/pybind11/archive/v2.4.3.tar.gz"],
+)
+
+new_local_repository(
+    name = "python_headers",
+    build_file = "@//bindings/python:python_headers.BUILD",
+    path = "/usr/include/python3.6",  # May be overwritten by setup.py.
+)
+

bindings/python/BUILD

@@ -0,0 +1,3 @@
+exports_files(glob(["*.BUILD"]))
+exports_files(["build_defs.bzl"])
+

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",
+    ],
+)
+

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)

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

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()

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()
+        }),
+    )

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"],
+)

bindings/python/python_headers.BUILD

@@ -0,0 +1,6 @@
+cc_library(
+    name = "python_headers",
+    hdrs = glob(["**/*.h"]),
+    includes = ["."],
+    visibility = ["//visibility:public"],
+)

bindings/python/requirements.txt

@@ -0,0 +1,2 @@
+absl-py>=0.7.1
+

setup.py

@@ -0,0 +1,124 @@
+import os
+import posixpath
+import re
+import shutil
+import sys
+
+from distutils import sysconfig
+import setuptools
+from setuptools.command import build_ext
+
+
+here = os.path.dirname(os.path.abspath(__file__))
+
+
+IS_WINDOWS = sys.platform.startswith('win')
+
+
+def _get_version():
+  """Parse the version string from __init__.py."""
+  with open(os.path.join(here, 'bindings', 'python', 'benchmark', '__init__.py')) as f:
+    try:
+      version_line = next(
+          line for line in f if line.startswith('__version__'))
+    except StopIteration:
+      raise ValueError('__version__ not defined in __init__.py')
+    else:
+      ns = {}
+      exec(version_line, ns)  # pylint: disable=exec-used
+      return ns['__version__']
+
+
+def _parse_requirements(path):
+  with open(os.path.join(here, path)) as f:
+    return [
+        line.rstrip() for line in f
+        if not (line.isspace() or line.startswith('#'))
+    ]
+
+
+class BazelExtension(setuptools.Extension):
+  """A C/C++ extension that is defined as a Bazel BUILD target."""
+
+  def __init__(self, name, bazel_target):
+    self.bazel_target = bazel_target
+    self.relpath, self.target_name = (
+        posixpath.relpath(bazel_target, '//').split(':'))
+    setuptools.Extension.__init__(self, name, sources=[])
+
+
+class BuildBazelExtension(build_ext.build_ext):
+  """A command that runs Bazel to build a C/C++ extension."""
+
+  def run(self):
+    for ext in self.extensions:
+      self.bazel_build(ext)
+    build_ext.build_ext.run(self)
+
+  def bazel_build(self, ext):
+    with open('WORKSPACE', 'r') as f:
+      workspace_contents = f.read()
+
+    with open('WORKSPACE', 'w') as f:
+      f.write(re.sub(
+          r'(?<=path = ").*(?=",  # May be overwritten by setup\.py\.)',
+          sysconfig.get_python_inc().replace(os.path.sep, posixpath.sep),
+          workspace_contents))
+
+    if not os.path.exists(self.build_temp):
+      os.makedirs(self.build_temp)
+
+    bazel_argv = [
+        'bazel',
+        'build',
+        ext.bazel_target,
+        '--symlink_prefix=' + os.path.join(self.build_temp, 'bazel-'),
+        '--compilation_mode=' + ('dbg' if self.debug else 'opt'),
+    ]
+
+    if IS_WINDOWS:
+      # Link with python*.lib.
+      for library_dir in self.library_dirs:
+        bazel_argv.append('--linkopt=/LIBPATH:' + library_dir)
+
+    self.spawn(bazel_argv)
+
+    shared_lib_suffix = '.dll' if IS_WINDOWS else '.so'
+    ext_bazel_bin_path = os.path.join(
+        self.build_temp, 'bazel-bin',
+        ext.relpath, ext.target_name + shared_lib_suffix)
+    ext_dest_path = self.get_ext_fullpath(ext.name)
+    ext_dest_dir = os.path.dirname(ext_dest_path)
+    if not os.path.exists(ext_dest_dir):
+      os.makedirs(ext_dest_dir)
+    shutil.copyfile(ext_bazel_bin_path, ext_dest_path)
+
+
+setuptools.setup(
+    name='google-benchmark',
+    version=_get_version(),
+    url='https://github.com/google/benchmark',
+    description='A library to benchmark code snippets.',
+    author='Google',
+    author_email='benchmark-py@google.com',
+    # Contained modules and scripts.
+    package_dir={'': 'bindings/python'},
+    packages=setuptools.find_packages('bindings/python'),
+    install_requires=_parse_requirements('bindings/python/requirements.txt'),
+    cmdclass=dict(build_ext=BuildBazelExtension),
+    ext_modules=[BazelExtension('benchmark._benchmark', '//bindings/python/benchmark:_benchmark')],
+    zip_safe=False,
+    # PyPI package information.
+    classifiers=[
+        'Development Status :: 4 - Beta',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
+        'Topic :: Software Development :: Testing',
+        'Topic :: System :: Benchmark',
+    ],
+    license='Apache 2.0',
+    keywords='benchmark',
+)