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