-
Notifications
You must be signed in to change notification settings - Fork 7
/
__init__.py
351 lines (293 loc) · 12.9 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
'''Python lisp-like tracing library. Prints to stdout a nested
display of function calls with arguments and return values. Also
prints exceptions when exceptions are thrown.
It works by temporarily replacing all functions/methods within the
listed classes/modules with traced versions. Then when the 'trace'
block exits, all the original values are restored.
Notes:
* When tracing classes, only the methods defined in that class
are traced, inherited methods are not traced.
* Tracing __repr__ will cause a stack overflow, since this method is
used to print out trace arguments. The tracer will always skip
tracing this method, even when using `include_hidden`.
* You can change where the trace goes by redefining `tracer`. It
should be a function that takes f, *args, **kwargs and calls f with
the args. By default it points to `trace`.
Usage:
with trace_on([Class1, module1, Class2, module2]):
module1.function1(arg1, arg2)
x = new Class1()
x.method1(arg1, arg2)
'''
from contextlib import contextmanager, closing
import inspect
import os
import sys
import itertools
indentchar = "| "
def _name(f):
'''Get an appropriate name for the object, for printing to the trace log'''
nattr = None
mattr = getattr(f, '__module__', None)
if mattr:
nattr = getattr(f, "__name__", None)
else:
# partials
fattr = getattr(f, 'func', None)
if fattr:
mattr = getattr(fattr, '__module__')
nattr = getattr(fattr, '__name__') + '__partial__'
if mattr and nattr:
return "%s.%s" % (mattr, nattr)
else:
return repr(f)
class Formatter(object):
def __init__(self):
pass
def format_input(self, level, f, args, kwargs):
return "%s- %s(%s)" % \
(level * indentchar, f,
", ".join(map(repr, args) +
map(lambda i: "%s=%s" % (i[0], repr(i[1])),
kwargs.items())))
def format_output(self, level, returnval, exception):
return "%s-> %s%s" % (level * indentchar,
"!!!" if exception else "",
repr(returnval))
def _get_function_mapping(o):
'''Returns a 2-tuple, with the first element being the object
belonging to the function/method that can be recognized from
the Frame info, the second element is an object that contains
info that can be printed to the trace log.
'''
# for regular functions, the identifier is the code object that appears in the frame
# and the function itself is where the tracing info lies
if inspect.ismethod(o) or inspect.isfunction(o):
if hasattr(o, 'func'):
i = o.func.func_code
elif hasattr(o, 'im_func'):
i = o.im_func.func_code
elif hasattr(o.__call__, 'im_func'):
i = o.__call__.im_func.func_code
elif hasattr(o, 'func_code'):
i = o.func_code
return (i, o)
# for objects that implement __call__, like MultiMethods, the identifier is the instance
# since each one is conceptually a different function. The tracing info is really just
# the module and the function's name is '__call__'. The first arg is what's important
# (the instance)
if hasattr(o, '__call__'):
# print "got mm %s" % o
try:
# return (o.__call__.im_func.func_code, o.__call__)
return (o.__call__, o.__call__)
except:
# print o
return (None, None)
return None
class Tracer(object):
def _get_functions(self, functions, depths):
'''sets some attributes:
functions = mapping of identifiers to functions
depths = mapping of identifiers to depth
'''
self.functions = {}
self.depths = {}
functions = set(functions) | set(depths.keys())
for f in functions:
ident, info_obj = _get_function_mapping(f)
self.functions[ident] = info_obj
if f in depths:
self.depths[ident] = depths[f]
# self.count = 0
# self.skipped = 0
self.min_depth = sys.maxint
# self.counts = {}
self.no_trace = set()
def __init__(self, functions, formatter=None, depths=None):
self.formatter = formatter or Formatter()
self.exception_frame = None
# keep our own call stack of just frames being traced
self.tracedframes = [] # tuples of (frame, maxdepth, is_traced)
self._get_functions(functions, depths or {})
def _get_id(self, frame):
'''
Given a frame, figure out what function/method is being called.
'''
f = frame.f_code
# self.counts[f] = self.counts.get(f, 0) + 1
if f in self.functions:
return f # if it's in the functions dict, we know it's correct
else:
# it could be an object that implements __call__. Find __call__.
args, varargs, keywords, localz = inspect.getargvalues(frame)
if args:
try:
# first arg is self, the instance
cf = localz[args[0]].__call__
if cf.im_func.func_code is f and cf in self.functions:
return cf
except BaseException:
pass
return None
def _min_depths(self):
'''Depth-controlled functions will limit the displayed call depth,
find the most restrictive one (the minimum depth)'''
depths = [fmd[1] for fmd in self.tracedframes] or [sys.maxint]
return min(depths)
@property
def level(self):
return len(self.tracedframes)
def _method_or_function_call(self, frame, ident):
f = self.functions[ident]
args = inspect.getargvalues(frame)
if inspect.ismethod(f):
locs = args.locals.copy()
f_self = locs.pop(args.args[0])
self.trace_in("%s.%s" % (repr(f_self), f.__name__), [], locs)
else:
# regular function
self.trace_in(_name(f), [], args.locals)
def tracefunc(self, frame, event, arg):
# self.counts[event] = self.counts.get(event, 0) + 1
try:
if event == 'call' and frame.f_code not in self.no_trace:
# return
# f = frame.f_code
# ident = f if f in self.functions else None
ident = self._get_id(frame)
if ident:
additional_depth = self.depths.get(ident, None)
min_depth_limit = self._min_depths()
if self.level < min_depth_limit:
if additional_depth is not None:
next_depth_limit = self.level + additional_depth
if next_depth_limit < min_depth_limit:
min_depth_limit = next_depth_limit
if self.level < min_depth_limit:
self._method_or_function_call(frame, ident)
if self.level <= min_depth_limit:
self.tracedframes.append((frame.f_back, min_depth_limit,
self.level < min_depth_limit))
if min_depth_limit < self.min_depth:
self.min_depth = min_depth_limit
else:
self.no_trace.add(frame.f_code)
if self.level >= self.min_depth:
# print "cut off! %s:%s" % (frame.f_code.co_filename, frame.f_lineno)
return None
elif event == 'return':
# print frame.f_code
if self.exception_frame:
self.exception_frame = None
elif self.tracedframes and self.tracedframes[-1][0] is frame.f_back:
# print self.tracedframes
if self.tracedframes[-1][2]:
# self.trace_out("%s: %s" % (self._get_id(frame), arg))
self.trace_out(arg)
_x, min_depth_limit, _y = self.tracedframes.pop()
if self.min_depth == min_depth_limit:
# recalculate min depth
self.min_depth = self._min_depths()
elif event == 'exception':
if self.tracedframes and self.tracedframes[-1][0] is frame.f_back:
# since both return and exception events get called for exceptions,
# save this frame so that we know it's the same trace entry when we get the
# return event.
self.exception_frame = frame
if self.tracedframes[-1][2]:
# self.trace_out("%s: %s" % (self._get_id(frame), arg))
self.trace_out(arg[0], exception=True)
_x, min_depth_limit, _y = self.tracedframes.pop()
if self.min_depth == min_depth_limit:
# recalculate min depth
self.min_depth = self._min_depths()
except:
# pass # just swallow errors to avoid interference with traced processes
raise # for debugging
return self.tracefunc
def close(self):
pass
# print "count=%s, skipped=%s" % (self.count, self.skipped)
# counts = sorted(self.counts.iteritems(), key=operator.itemgetter(1))
# counts = counts[-50:]
# print "count=%s, skipped=%s, counts=%s" % (self.count, self.skipped, counts)
class StdoutTracer(Tracer):
'''Print trace to stdout'''
def __init__(self, functions, formatter=None, depths=None):
super(StdoutTracer, self).__init__(functions, formatter=formatter, depths=depths)
def trace_in(self, f, args, kwargs):
print self.formatter.format_input(self.level, f, args, kwargs)
sys.stdout.flush()
def trace_out(self, r, exception=False):
print self.formatter.format_output(self.level - 1, r, exception)
sys.stdout.flush()
def close(self):
# print "closing " + str(self.outputfile)
# print "count=%s, skipped=%s, counts=%s" % (self.count, self.skipped, self.counts)
pass
class PerThreadFileTracer(Tracer):
'''Print trace to a file. To get thread safety, use a different
instance of this tracer for each thread.'''
def __init__(self, functions, formatter=None, depths=None, filename=None):
super(PerThreadFileTracer, self).__init__(functions, formatter=formatter, depths=depths)
d = os.path.dirname(filename)
if not os.path.exists(d):
os.makedirs(d)
# keep file we're writing to outside the state of this instance
# prevents replaced functions from trying to write to the wrong file
self.outputfile = open(filename, 'w')
def trace_in(self, f, args, kwargs):
# print "in %s %s %s %s" % (str(self.outputfile), f, args, kwargs)
self.outputfile.write(self.formatter.format_input(self.level, f, args, kwargs) + "\n")
def trace_out(self, r, exception=False):
self.outputfile.write(self.formatter.format_output(self.level - 1, r, exception) + "\n")
def close(self):
# print "closing " + str(self.outputfile)
# print "count=%s, skipped=%s, counts=%s" % (self.count, self.skipped, self.counts)
self.outputfile.close()
def _defined_this_module(parent, child):
'''Returns true if f is defined in the module o (or true if o is a class)'''
if inspect.ismodule(parent):
return parent.__name__ == getattr(child, '__module__', None)
return True
def mapcat(f, lst):
return list(itertools.chain.from_iterable(map(f, lst)))
def all(o, include_hidden=False):
'''Return all the functions/methods in the given object.'''
def r(x, y):
'''Return list of functions x plus all the function members of y'''
n, v = y # getmembers returns name/value tuples
if not n.startswith("__")\
and (include_hidden or
not (include_hidden or n.startswith("_")))\
and _defined_this_module(o, v):
if inspect.isclass(v):
# print "adding class %s" % v
x.extend(all(v))
return x
else:
# print "adding function %s " % v
x.append(v)
return x
else:
return list(x)
return reduce(r, inspect.getmembers(o, callable), [])
def add_all_at_depth(dct, module, lvl):
'''takes a depth dict, module, and level, and adds all the functions
in the module to the depth dict at the given level.
Returns: the dict with new values
'''
fns = all(module)
for f in fns:
dct[f] = lvl
return dct
@contextmanager
def trace_on(objs=None, tracer=None):
tracer = tracer or StdoutTracer(objs)
sys.settrace(tracer.tracefunc)
with closing(tracer):
try:
yield
finally:
sys.settrace(None)