|
1 import pstats |
|
2 import cPickle |
|
3 import zlib |
|
4 |
|
5 class PickleStats(object): |
|
6 def __init__(self, stats): |
|
7 self.stats = stats |
|
8 |
|
9 def create_stats(self): |
|
10 "pstats.Stats checks for the existence of this method to see if it can load from an object" |
|
11 pass |
|
12 |
|
13 def from_file(fileobj): |
|
14 "load ppstats from an open file object" |
|
15 stats = cPickle.load(fileobj) |
|
16 ps = PickleStats(stats) |
|
17 return Stats(ps) |
|
18 |
|
19 def from_filename(filename): |
|
20 "load ppstats from a filename" |
|
21 fileobj = open(filename, 'rb') |
|
22 return from_file(fileobj) |
|
23 |
|
24 def from_gz_file(fileobj): |
|
25 "load ppstats from an open file containing gzipped data" |
|
26 data = fileobj.read() |
|
27 stats = cPickle.loads(zlib.decompress(data)) |
|
28 ps = PickleStats(stats) |
|
29 return Stats(ps) |
|
30 |
|
31 def from_gz_filename(filename): |
|
32 "load ppstats from a file containing gzipped data, by filename" |
|
33 fileobj = open(filename, 'rb') |
|
34 return from_gz_file(fileobj) |
|
35 |
|
36 def from_gz(gz_string): |
|
37 "load ppstats from a string of gzipped data" |
|
38 return Stats(PickleStats(cPickle.loads(zlib.decompress(gz_string)))) |
|
39 |
|
40 def from_stats(stats): |
|
41 "load ppstats from a stats object" |
|
42 return Stats(PickleStats(stats)) |
|
43 |
|
44 def from_string(stats_string): |
|
45 return Stats(PickleStats(cPickle.loads(stats_string))) |
|
46 |
|
47 class Stats(pstats.Stats): |
|
48 def __init__(self, *args, **kwargs): |
|
49 pstats.Stats.__init__(self, *args) |
|
50 self.replace_dirs = {} |
|
51 |
|
52 def set_output(self, stream): |
|
53 "redirect output of print_stats to the file object <stream>" |
|
54 self.stream = stream |
|
55 |
|
56 def hide_directory(self, dirname, replacement=''): |
|
57 "replace occurences of <dirname> in filenames with <replacement>" |
|
58 self.replace_dirs[dirname] = replacement |
|
59 |
|
60 def func_strip_path(self, func_name): |
|
61 "take a filename, line, name tuple and mangle appropiately" |
|
62 filename, line, name = func_name |
|
63 |
|
64 for dirname in self.replace_dirs: |
|
65 filename = filename.replace(dirname, self.replace_dirs[dirname]) |
|
66 |
|
67 return filename, line, name |
|
68 |
|
69 def strip_dirs(self): |
|
70 "strip irrelevant/redundant directories from filenames in profile data" |
|
71 func_std_string = pstats.func_std_string |
|
72 |
|
73 oldstats = self.stats |
|
74 self.stats = newstats = {} |
|
75 max_name_len = 0 |
|
76 for func, (cc, nc, tt, ct, callers) in oldstats.iteritems(): |
|
77 newfunc = self.func_strip_path(func) |
|
78 if len(func_std_string(newfunc)) > max_name_len: |
|
79 max_name_len = len(func_std_string(newfunc)) |
|
80 newcallers = {} |
|
81 for func2, caller in callers.iteritems(): |
|
82 newcallers[self.func_strip_path(func2)] = caller |
|
83 |
|
84 if newfunc in newstats: |
|
85 newstats[newfunc] = add_func_stats( |
|
86 newstats[newfunc], |
|
87 (cc, nc, tt, ct, newcallers)) |
|
88 else: |
|
89 newstats[newfunc] = (cc, nc, tt, ct, newcallers) |
|
90 old_top = self.top_level |
|
91 self.top_level = new_top = {} |
|
92 for func in old_top: |
|
93 new_top[self.func_strip_path(func)] = None |
|
94 |
|
95 self.max_name_len = max_name_len |
|
96 |
|
97 self.fcn_list = None |
|
98 self.all_callees = None |
|
99 return self |
|
100 |
|
101 def dump_stats_pickle(self): |
|
102 "return a string containing picked stats information (dump_stats uses marshall)" |
|
103 return cPickle.dumps(self.stats) |
|
104 |
|
105 def load_stats_pickle(self, pickle_string): |
|
106 "load from string returned by dump_stats_pickle" |
|
107 return self.load_stats(PickleStats(cPickle.load(pickle_string))) |