__init__.py 4.3 KB

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