@@ -90,6 +90,17 @@ def reindent(src, indent):
90
90
"""Helper to reindent a multi-line statement."""
91
91
return src .replace ("\n " , "\n " + " " * indent )
92
92
93
+ def _template_func (setup , func ):
94
+ """Create a timer function. Used if the "statement" is a callable."""
95
+ def inner (_it , _timer ):
96
+ setup ()
97
+ _t0 = _timer ()
98
+ for _i in _it :
99
+ func ()
100
+ _t1 = _timer ()
101
+ return _t1 - _t0
102
+ return inner
103
+
93
104
class Timer :
94
105
"""Class for timing execution speed of small code snippets.
95
106
@@ -109,14 +120,32 @@ class Timer:
109
120
def __init__ (self , stmt = "pass" , setup = "pass" , timer = default_timer ):
110
121
"""Constructor. See class doc string."""
111
122
self .timer = timer
112
- stmt = reindent (stmt , 8 )
113
- setup = reindent (setup , 4 )
114
- src = template % {'stmt' : stmt , 'setup' : setup }
115
- self .src = src # Save for traceback display
116
- code = compile (src , dummy_src_name , "exec" )
117
123
ns = {}
118
- exec code in globals (), ns
119
- self .inner = ns ["inner" ]
124
+ if isinstance (stmt , basestring ):
125
+ stmt = reindent (stmt , 8 )
126
+ if isinstance (setup , basestring ):
127
+ setup = reindent (setup , 4 )
128
+ src = template % {'stmt' : stmt , 'setup' : setup }
129
+ elif callable (setup ):
130
+ src = template % {'stmt' : stmt , 'setup' : '_setup()' }
131
+ ns ['_setup' ] = setup
132
+ else :
133
+ raise ValueError ("setup is neither a string nor callable" )
134
+ self .src = src # Save for traceback display
135
+ code = compile (src , dummy_src_name , "exec" )
136
+ exec code in globals (), ns
137
+ self .inner = ns ["inner" ]
138
+ elif callable (stmt ):
139
+ self .src = None
140
+ if isinstance (setup , basestring ):
141
+ _setup = setup
142
+ def setup ():
143
+ exec _setup in globals (), ns
144
+ elif not callable (setup ):
145
+ raise ValueError ("setup is neither a string nor callable" )
146
+ self .inner = _template_func (setup , stmt )
147
+ else :
148
+ raise ValueError ("stmt is neither a string nor callable" )
120
149
121
150
def print_exc (self , file = None ):
122
151
"""Helper to print a traceback from the timed code.
@@ -136,10 +165,13 @@ def print_exc(self, file=None):
136
165
sent; it defaults to sys.stderr.
137
166
"""
138
167
import linecache , traceback
139
- linecache .cache [dummy_src_name ] = (len (self .src ),
140
- None ,
141
- self .src .split ("\n " ),
142
- dummy_src_name )
168
+ if self .src is not None :
169
+ linecache .cache [dummy_src_name ] = (len (self .src ),
170
+ None ,
171
+ self .src .split ("\n " ),
172
+ dummy_src_name )
173
+ # else the source is already stored somewhere else
174
+
143
175
traceback .print_exc (file = file )
144
176
145
177
def timeit (self , number = default_number ):
@@ -189,6 +221,16 @@ def repeat(self, repeat=default_repeat, number=default_number):
189
221
r .append (t )
190
222
return r
191
223
224
+ def timeit (stmt = "pass" , setup = "pass" , timer = default_timer ,
225
+ number = default_number ):
226
+ """Convenience function to create Timer object and call timeit method."""
227
+ return Timer (stmt , setup , timer ).timeit (number )
228
+
229
+ def repeat (stmt = "pass" , setup = "pass" , timer = default_timer ,
230
+ repeat = default_repeat , number = default_number ):
231
+ """Convenience function to create Timer object and call repeat method."""
232
+ return Timer (stmt , setup , timer ).repeat (repeat , number )
233
+
192
234
def main (args = None ):
193
235
"""Main program, used when run as a script.
194
236
0 commit comments