third-party/benchmark/bindings/python/google_benchmark/__init__.py

   1 # Copyright 2020 Google Inc. All rights reserved.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 """Python benchmarking utilities.
  15
  16 Example usage:
  17   import google_benchmark as benchmark
  18
  19   @benchmark.register
  20   def my_benchmark(state):
  21       ...  # Code executed outside `while` loop is not timed.
  22
  23       while state:
  24         ...  # Code executed within `while` loop is timed.
  25
  26   if __name__ == '__main__':
  27     benchmark.main()
  28 """
  29
  30 from absl import app
  31 from google_benchmark import _benchmark
  32 from google_benchmark._benchmark import (
  33     Counter,
  34     kNanosecond,
  35     kMicrosecond,
  36     kMillisecond,
  37     kSecond,
  38     oNone,
  39     o1,
  40     oN,
  41     oNSquared,
  42     oNCubed,
  43     oLogN,
  44     oNLogN,
  45     oAuto,
  46     oLambda,
  47 )
  48
  49
  50 __all__ = [
  51     "register",
  52     "main",
  53     "Counter",
  54     "kNanosecond",
  55     "kMicrosecond",
  56     "kMillisecond",
  57     "kSecond",
  58     "oNone",
  59     "o1",
  60     "oN",
  61     "oNSquared",
  62     "oNCubed",
  63     "oLogN",
  64     "oNLogN",
  65     "oAuto",
  66     "oLambda",
  67 ]
  68
  69 __version__ = "0.2.0"
  70
  71
  72 class __OptionMaker:
  73     """A stateless class to collect benchmark options.
  74
  75     Collect all decorator calls like @option.range(start=0, limit=1<<5).
  76     """
  77
  78     class Options:
  79         """Pure data class to store options calls, along with the benchmarked function."""
  80
  81         def __init__(self, func):
  82             self.func = func
  83             self.builder_calls = []
  84
  85     @classmethod
  86     def make(cls, func_or_options):
  87         """Make Options from Options or the benchmarked function."""
  88         if isinstance(func_or_options, cls.Options):
  89             return func_or_options
  90         return cls.Options(func_or_options)
  91
  92     def __getattr__(self, builder_name):
  93         """Append option call in the Options."""
  94
  95         # The function that get returned on @option.range(start=0, limit=1<<5).
  96         def __builder_method(*args, **kwargs):
  97
  98             # The decorator that get called, either with the benchmared function
  99             # or the previous Options
 100             def __decorator(func_or_options):
 101                 options = self.make(func_or_options)
 102                 options.builder_calls.append((builder_name, args, kwargs))
 103                 # The decorator returns Options so it is not technically a decorator
 104                 # and needs a final call to @regiser
 105                 return options
 106
 107             return __decorator
 108
 109         return __builder_method
 110
 111
 112 # Alias for nicer API.
 113 # We have to instantiate an object, even if stateless, to be able to use __getattr__
 114 # on option.range
 115 option = __OptionMaker()
 116
 117
 118 def register(undefined=None, *, name=None):
 119     """Register function for benchmarking."""
 120     if undefined is None:
 121         # Decorator is called without parenthesis so we return a decorator
 122         return lambda f: register(f, name=name)
 123
 124     # We have either the function to benchmark (simple case) or an instance of Options
 125     # (@option._ case).
 126     options = __OptionMaker.make(undefined)
 127
 128     if name is None:
 129         name = options.func.__name__
 130
 131     # We register the benchmark and reproduce all the @option._ calls onto the
 132     # benchmark builder pattern
 133     benchmark = _benchmark.RegisterBenchmark(name, options.func)
 134     for name, args, kwargs in options.builder_calls[::-1]:
 135         getattr(benchmark, name)(*args, **kwargs)
 136
 137     # return the benchmarked function because the decorator does not modify it
 138     return options.func
 139
 140
 141 def _flags_parser(argv):
 142     argv = _benchmark.Initialize(argv)
 143     return app.parse_flags_with_usage(argv)
 144
 145
 146 def _run_benchmarks(argv):
 147     if len(argv) > 1:
 148         raise app.UsageError("Too many command-line arguments.")
 149     return _benchmark.RunSpecifiedBenchmarks()
 150
 151
 152 def main(argv=None):
 153     return app.run(_run_benchmarks, argv=argv, flags_parser=_flags_parser)
 154
 155
 156 # Methods for use with custom main function.
 157 initialize = _benchmark.Initialize
 158 run_benchmarks = _benchmark.RunSpecifiedBenchmarks