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 import atexit
  30
  31 from absl import app
  32
  33 from google_benchmark import _benchmark
  34 from google_benchmark._benchmark import (
  35     Counter as Counter,
  36     State as State,
  37     kMicrosecond as kMicrosecond,
  38     kMillisecond as kMillisecond,
  39     kNanosecond as kNanosecond,
  40     kSecond as kSecond,
  41     o1 as o1,
  42     oAuto as oAuto,
  43     oLambda as oLambda,
  44     oLogN as oLogN,
  45     oN as oN,
  46     oNCubed as oNCubed,
  47     oNLogN as oNLogN,
  48     oNone as oNone,
  49     oNSquared as oNSquared,
  50 )
  51 from google_benchmark.version import __version__ as __version__
  52
  53
  54 class __OptionMaker:
  55     """A stateless class to collect benchmark options.
  56
  57     Collect all decorator calls like @option.range(start=0, limit=1<<5).
  58     """
  59
  60     class Options:
  61         """Pure data class to store options calls, along with the benchmarked function."""
  62
  63         def __init__(self, func):
  64             self.func = func
  65             self.builder_calls = []
  66
  67     @classmethod
  68     def make(cls, func_or_options):
  69         """Make Options from Options or the benchmarked function."""
  70         if isinstance(func_or_options, cls.Options):
  71             return func_or_options
  72         return cls.Options(func_or_options)
  73
  74     def __getattr__(self, builder_name):
  75         """Append option call in the Options."""
  76
  77         # The function that get returned on @option.range(start=0, limit=1<<5).
  78         def __builder_method(*args, **kwargs):
  79             # The decorator that get called, either with the benchmared function
  80             # or the previous Options
  81             def __decorator(func_or_options):
  82                 options = self.make(func_or_options)
  83                 options.builder_calls.append((builder_name, args, kwargs))
  84                 # The decorator returns Options so it is not technically a decorator
  85                 # and needs a final call to @register
  86                 return options
  87
  88             return __decorator
  89
  90         return __builder_method
  91
  92
  93 # Alias for nicer API.
  94 # We have to instantiate an object, even if stateless, to be able to use __getattr__
  95 # on option.range
  96 option = __OptionMaker()
  97
  98
  99 def register(undefined=None, *, name=None):
 100     """Register function for benchmarking."""
 101     if undefined is None:
 102         # Decorator is called without parenthesis so we return a decorator
 103         return lambda f: register(f, name=name)
 104
 105     # We have either the function to benchmark (simple case) or an instance of Options
 106     # (@option._ case).
 107     options = __OptionMaker.make(undefined)
 108
 109     if name is None:
 110         name = options.func.__name__
 111
 112     # We register the benchmark and reproduce all the @option._ calls onto the
 113     # benchmark builder pattern
 114     benchmark = _benchmark.RegisterBenchmark(name, options.func)
 115     for name, args, kwargs in options.builder_calls[::-1]:
 116         getattr(benchmark, name)(*args, **kwargs)
 117
 118     # return the benchmarked function because the decorator does not modify it
 119     return options.func
 120
 121
 122 def _flags_parser(argv):
 123     argv = _benchmark.Initialize(argv)
 124     return app.parse_flags_with_usage(argv)
 125
 126
 127 def _run_benchmarks(argv):
 128     if len(argv) > 1:
 129         raise app.UsageError("Too many command-line arguments.")
 130     return _benchmark.RunSpecifiedBenchmarks()
 131
 132
 133 def main(argv=None):
 134     return app.run(_run_benchmarks, argv=argv, flags_parser=_flags_parser)
 135
 136
 137 # Methods for use with custom main function.
 138 initialize = _benchmark.Initialize
 139 run_benchmarks = _benchmark.RunSpecifiedBenchmarks
 140 atexit.register(_benchmark.ClearRegisteredBenchmarks)