benchmarking.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221
  1. """benchmarking through py.test"""
  2. import py
  3. from py.__.test.item import Item
  4. from py.__.test.terminal.terminal import TerminalSession
  5. from math import ceil as _ceil, floor as _floor, log10
  6. import timeit
  7. from inspect import getsource
  8. # from IPython.Magic.magic_timeit
  9. units = ["s", "ms", "us", "ns"]
  10. scaling = [1, 1e3, 1e6, 1e9]
  11. unitn = {s: i for i, s in enumerate(units)}
  12. precision = 3
  13. # like py.test Directory but scan for 'bench_<smth>.py'
  14. class Directory(py.test.collect.Directory):
  15. def filefilter(self, path):
  16. b = path.purebasename
  17. ext = path.ext
  18. return b.startswith('bench_') and ext == '.py'
  19. # like py.test Module but scane for 'bench_<smth>' and 'timeit_<smth>'
  20. class Module(py.test.collect.Module):
  21. def funcnamefilter(self, name):
  22. return name.startswith('bench_') or name.startswith('timeit_')
  23. # Function level benchmarking driver
  24. class Timer(timeit.Timer):
  25. def __init__(self, stmt, setup='pass', timer=timeit.default_timer, globals=globals()):
  26. # copy of timeit.Timer.__init__
  27. # similarity index 95%
  28. self.timer = timer
  29. stmt = timeit.reindent(stmt, 8)
  30. setup = timeit.reindent(setup, 4)
  31. src = timeit.template % {'stmt': stmt, 'setup': setup}
  32. self.src = src # Save for traceback display
  33. code = compile(src, timeit.dummy_src_name, "exec")
  34. ns = {}
  35. #exec(code, globals(), ns) -- original timeit code
  36. exec(code, globals, ns) # -- we use caller-provided globals instead
  37. self.inner = ns["inner"]
  38. class Function(py.__.test.item.Function):
  39. def __init__(self, *args, **kw):
  40. super().__init__(*args, **kw)
  41. self.benchtime = None
  42. self.benchtitle = None
  43. def execute(self, target, *args):
  44. # get func source without first 'def func(...):' line
  45. src = getsource(target)
  46. src = '\n'.join( src.splitlines()[1:] )
  47. # extract benchmark title
  48. if target.func_doc is not None:
  49. self.benchtitle = target.func_doc
  50. else:
  51. self.benchtitle = src.splitlines()[0].strip()
  52. # XXX we ignore args
  53. timer = Timer(src, globals=target.func_globals)
  54. if self.name.startswith('timeit_'):
  55. # from IPython.Magic.magic_timeit
  56. repeat = 3
  57. number = 1
  58. for i in range(1, 10):
  59. t = timer.timeit(number)
  60. if t >= 0.2:
  61. number *= (0.2 / t)
  62. number = int(_ceil(number))
  63. break
  64. if t <= 0.02:
  65. # we are not close enough to that 0.2s
  66. number *= 10
  67. else:
  68. # since we are very close to be > 0.2s we'd better adjust number
  69. # so that timing time is not too high
  70. number *= (0.2 / t)
  71. number = int(_ceil(number))
  72. break
  73. self.benchtime = min(timer.repeat(repeat, number)) / number
  74. # 'bench_<smth>'
  75. else:
  76. self.benchtime = timer.timeit(1)
  77. class BenchSession(TerminalSession):
  78. def header(self, colitems):
  79. super().header(colitems)
  80. def footer(self, colitems):
  81. super().footer(colitems)
  82. self.out.write('\n')
  83. self.print_bench_results()
  84. def print_bench_results(self):
  85. self.out.write('==============================\n')
  86. self.out.write(' *** BENCHMARKING RESULTS *** \n')
  87. self.out.write('==============================\n')
  88. self.out.write('\n')
  89. # benchname, time, benchtitle
  90. results = []
  91. for item, outcome in self._memo:
  92. if isinstance(item, Item):
  93. best = item.benchtime
  94. if best is None:
  95. # skipped or failed benchmarks
  96. tstr = '---'
  97. else:
  98. # from IPython.Magic.magic_timeit
  99. if best > 0.0:
  100. order = min(-int(_floor(log10(best)) // 3), 3)
  101. else:
  102. order = 3
  103. tstr = "%.*g %s" % (
  104. precision, best * scaling[order], units[order])
  105. results.append( [item.name, tstr, item.benchtitle] )
  106. # dot/unit align second column
  107. # FIXME simpler? this is crappy -- shame on me...
  108. wm = [0]*len(units)
  109. we = [0]*len(units)
  110. for s in results:
  111. tstr = s[1]
  112. n, u = tstr.split()
  113. # unit n
  114. un = unitn[u]
  115. try:
  116. m, e = n.split('.')
  117. except ValueError:
  118. m, e = n, ''
  119. wm[un] = max(len(m), wm[un])
  120. we[un] = max(len(e), we[un])
  121. for s in results:
  122. tstr = s[1]
  123. n, u = tstr.split()
  124. un = unitn[u]
  125. try:
  126. m, e = n.split('.')
  127. except ValueError:
  128. m, e = n, ''
  129. m = m.rjust(wm[un])
  130. e = e.ljust(we[un])
  131. if e.strip():
  132. n = '.'.join((m, e))
  133. else:
  134. n = ' '.join((m, e))
  135. # let's put the number into the right place
  136. txt = ''
  137. for i in range(len(units)):
  138. if i == un:
  139. txt += n
  140. else:
  141. txt += ' '*(wm[i] + we[i] + 1)
  142. s[1] = '%s %s' % (txt, u)
  143. # align all columns besides the last one
  144. for i in range(2):
  145. w = max(len(s[i]) for s in results)
  146. for s in results:
  147. s[i] = s[i].ljust(w)
  148. # show results
  149. for s in results:
  150. self.out.write('%s | %s | %s\n' % tuple(s))
  151. def main(args=None):
  152. # hook our Directory/Module/Function as defaults
  153. from py.__.test import defaultconftest
  154. defaultconftest.Directory = Directory
  155. defaultconftest.Module = Module
  156. defaultconftest.Function = Function
  157. # hook BenchSession as py.test session
  158. config = py.test.config
  159. config._getsessionclass = lambda: BenchSession
  160. py.test.cmdline.main(args)