4444
4545
4646_HRULE = '-' .join (['' for i in range (80 )])
47- ATTR_BENCHMARK = '__benchmark__'
47+
48+ #: this function is used to pre-process the arguments as expected by the __benchmark__ and __setup__ entry points
4849ATTR_PROCESS_ARGS = '__process_args__'
50+ #: gets called with the preprocessed arguments before __benchmark__
51+ ATTR_SETUP = '__setup__'
52+ #: gets called with the preprocessed arguments N times
53+ ATTR_BENCHMARK = '__benchmark__'
54+ #: performs any teardown needed in the benchmark
55+ ATTR_TEARDOWN = '__teardown__'
4956
5057
5158def ccompile (name , code ):
@@ -131,7 +138,6 @@ def _call_attr(self, attr_name, *args):
131138 return attr (* args )
132139
133140 def run (self ):
134- print (_HRULE )
135141 if self ._run_once :
136142 print ("### %s, exactly one iteration (no warmup curves)" % (self .bench_module .__name__ ))
137143 else :
@@ -143,10 +149,15 @@ def run(self):
143149 # default args processor considers all args as ints
144150 args = list (map (int , self .bench_args ))
145151
146- print ("### args = %s" % args )
152+ print ("### args = " , args )
147153 print (_HRULE )
148154
155+ print ("### setup ... " )
156+ self ._call_attr (ATTR_SETUP , * args )
157+ print ("### start benchmark ... " )
158+
149159 bench_func = self ._get_attr (ATTR_BENCHMARK )
160+ durations = []
150161 if bench_func and hasattr (bench_func , '__call__' ):
151162 if self .warmup :
152163 print ("### warming up for %s iterations ... " % self .warmup )
@@ -156,18 +167,31 @@ def run(self):
156167 for iteration in range (self .iterations ):
157168 start = time ()
158169 bench_func (* args )
159- duration = "%.3f" % (time () - start )
170+ duration = time () - start
171+ durations .append (duration )
172+ duration_str = "%.3f" % duration
160173 if self ._run_once :
161- print ("@@@ name=%s, duration=%s" % (self .bench_module .__name__ , duration ))
174+ print ("@@@ name=%s, duration=%s" % (self .bench_module .__name__ , duration_str ))
162175 else :
163- print ("### iteration=%s, name=%s, duration=%s" % (iteration , self .bench_module .__name__ , duration ))
176+ print ("### iteration=%s, name=%s, duration=%s" % (iteration , self .bench_module .__name__ , duration_str ))
177+
178+ print (_HRULE )
179+ print ("### teardown ... " )
180+ self ._call_attr (ATTR_TEARDOWN )
181+ print ("### benchmark complete" )
182+ print (_HRULE )
183+ print ("### BEST duration: %.3f s" % min (durations ))
184+ print ("### WORST duration: %.3f s" % max (durations ))
185+ print ("### AVG duration: %.3f" % (sum (durations ) / len (durations )))
186+ print (_HRULE )
164187
165188
166- def run_benchmark (prog , args ):
189+ def run_benchmark (args ):
167190 warmup = 0
168191 iterations = 1
169192 bench_file = None
170193 bench_args = []
194+ paths = []
171195
172196 i = 0
173197 while i < len (args ):
@@ -177,19 +201,36 @@ def run_benchmark(prog, args):
177201 iterations = _as_int (args [i ])
178202 elif arg .startswith ("--iterations" ):
179203 iterations = _as_int (arg .split ("=" )[1 ])
204+
180205 elif arg == '-w' :
181206 i += 1
182207 warmup = _as_int (args [i ])
183208 elif arg .startswith ("--warmup" ):
184209 warmup = _as_int (arg .split ("=" )[1 ])
210+
211+ elif arg == '-p' :
212+ i += 1
213+ paths = args [i ].split ("," )
214+ elif arg .startswith ("--path" ):
215+ paths = arg .split ("=" )[1 ].split ("," )
216+
185217 elif bench_file is None :
186218 bench_file = arg
187219 else :
188220 bench_args .append (arg )
189221 i += 1
190222
223+ # set the paths if specified
224+ print (_HRULE )
225+ if paths :
226+ for pth in paths :
227+ print ("### adding module path: %s" % pth )
228+ sys .path .append (pth )
229+ else :
230+ print ("### no extra module search paths specified" )
231+
191232 BenchRunner (bench_file , bench_args = bench_args , iterations = iterations , warmup = warmup ).run ()
192233
193234
194235if __name__ == '__main__' :
195- run_benchmark (sys .argv [0 ], sys . argv [ 1 :])
236+ run_benchmark (sys .argv [1 :])
0 commit comments