• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
16Example 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
30from absl import app
31from google_benchmark import _benchmark
32from 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__ = "1.6.1"
70
71
72class __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
115option = __OptionMaker()
116
117
118def 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
141def _flags_parser(argv):
142    argv = _benchmark.Initialize(argv)
143    return app.parse_flags_with_usage(argv)
144
145
146def _run_benchmarks(argv):
147    if len(argv) > 1:
148        raise app.UsageError("Too many command-line arguments.")
149    return _benchmark.RunSpecifiedBenchmarks()
150
151
152def 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.
157initialize = _benchmark.Initialize
158run_benchmarks = _benchmark.RunSpecifiedBenchmarks
159