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)))