src/pycallgraph.py

   1 """
   2 pycallgraph
   3
   4 U{http://pycallgraph.slowchop.com/}
   5
   6 Copyright Gerald Kaszuba 2007
   7
   8 This program is free software; you can redistribute it and/or modify
   9 it under the terms of the GNU General Public License as published by
  10 the Free Software Foundation; either version 2 of the License, or
  11 (at your option) any later version.
  12
  13 This program is distributed in the hope that it will be useful,
  14 but WITHOUT ANY WARRANTY; without even the implied warranty of
  15 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16 GNU General Public License for more details.
  17
  18 You should have received a copy of the GNU General Public License
  19 along with this program; if not, write to the Free Software
  20 Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
  21 """
  22
  23 __version__ = '0.4.1'
  24 __author__ = 'Gerald Kaszuba'
  25
  26 import inspect
  27 import sys
  28 import os
  29 import re
  30 import tempfile
  31 import time
  32 from distutils import sysconfig
  33
  34 # Initialise module variables.
  35 # TODO Move these into settings
  36 trace_filter = None
  37 time_filter = None
  38
  39
  40 def colourize_node(calls, total_time):
  41     value = float(total_time * 2 + calls) / 3
  42     return '%f %f %f' % (value / 2 + .5, value, 0.9)
  43
  44
  45 def colourize_edge(calls, total_time):
  46     value = float(total_time * 2 + calls) / 3
  47     return '%f %f %f' % (value / 2 + .5, value, 0.7)
  48
  49
  50 def reset_settings():
  51     global settings
  52     global graph_attributes
  53     global __version__
  54
  55     settings = {
  56         'node_attributes': {
  57            'label': r'%(func)s\ncalls: %(hits)i\ntotal time: %(total_time)f',
  58            'color': '%(col)s',
  59         },
  60         'node_colour': colourize_node,
  61         'edge_colour': colourize_edge,
  62         'dont_exclude_anything': False,
  63         'include_stdlib': True,
  64     }
  65
  66     # TODO: Move this into settings
  67     graph_attributes = {
  68         'graph': {
  69             'fontname': 'Verdana',
  70             'fontsize': 7,
  71             'fontcolor': '0 0 0.5',
  72             'label': r'Generated by Python Call Graph v%s\n' \
  73                 r'http://pycallgraph.slowchop.com' % __version__,
  74         },
  75         'node': {
  76             'fontname': 'Verdana',
  77             'fontsize': 7,
  78             'color': '.5 0 .9',
  79             'style': 'filled',
  80             'shape': 'rect',
  81         },
  82     }
  83
  84
  85 def reset_trace():
  86     """Resets all collected statistics. This is run automatically by
  87     start_trace(reset=True) and when the module is loaded.
  88     """
  89     global call_dict
  90     global call_stack
  91     global func_count
  92     global func_count_max
  93     global func_time
  94     global func_time_max
  95     global call_stack_timer
  96
  97     call_dict = {}
  98
  99     # current call stack
 100     call_stack = ['__main__']
 101
 102     # counters for each function
 103     func_count = {}
 104     func_count_max = 0
 105
 106     # accumative time per function
 107     func_time = {}
 108     func_time_max = 0
 109
 110     # keeps track of the start time of each call on the stack
 111     call_stack_timer = []
 112
 113
 114 class PyCallGraphException(Exception):
 115     """Exception used for pycallgraph"""
 116     pass
 117
 118
 119 class GlobbingFilter(object):
 120     """Filter module names using a set of globs.
 121
 122     Objects are matched against the exclude list first, then the include list.
 123     Anything that passes through without matching either, is excluded.
 124     """
 125
 126     def __init__(self, include=None, exclude=None, max_depth=None,
 127                  min_depth=None):
 128         if include is None and exclude is None:
 129             include = ['*']
 130             exclude = []
 131         elif include is None:
 132             include = ['*']
 133         elif exclude is None:
 134             exclude = []
 135         self.include = include
 136         self.exclude = exclude
 137         self.max_depth = max_depth or 9999
 138         self.min_depth = min_depth or 0
 139
 140     def __call__(self, stack, module_name=None, class_name=None,
 141                  func_name=None, full_name=None):
 142         from fnmatch import fnmatch
 143         if len(stack) > self.max_depth:
 144             return False
 145         if len(stack) < self.min_depth:
 146             return False
 147         for pattern in self.exclude:
 148             if fnmatch(full_name, pattern):
 149                 return False
 150         for pattern in self.include:
 151             if fnmatch(full_name, pattern):
 152                 return True
 153         return False
 154
 155
 156 def is_module_stdlib(file_name):
 157     """Returns True if the file_name is in the lib directory."""
 158     # TODO: Move these calls away from this function so it doesn't have to run
 159     # every time.
 160     lib_path = sysconfig.get_python_lib()
 161     path = os.path.split(lib_path)
 162     if path[1] == 'site-packages':
 163         lib_path = path[0]
 164     return file_name.lower().startswith(lib_path.lower())
 165
 166
 167 def start_trace(reset=True, filter_func=None, time_filter_func=None):
 168     """Begins a trace. Setting reset to True will reset all previously recorded
 169     trace data. filter_func needs to point to a callable function that accepts
 170     the parameters (call_stack, module_name, class_name, func_name, full_name).
 171     Every call will be passed into this function and it is up to the function
 172     to decide if it should be included or not. Returning False means the call
 173     will be filtered out and not included in the call graph.
 174     """
 175     global trace_filter
 176     global time_filter
 177     if reset:
 178         reset_trace()
 179
 180     if filter_func:
 181         trace_filter = filter_func
 182     else:
 183         trace_filter = GlobbingFilter(exclude=['pycallgraph.*'])
 184
 185     if time_filter_func:
 186         time_filter = time_filter_func
 187     else:
 188         time_filter = GlobbingFilter()
 189
 190     sys.settrace(tracer)
 191
 192
 193 def stop_trace():
 194     """Stops the currently running trace, if any."""
 195     sys.settrace(None)
 196
 197
 198 def tracer(frame, event, arg):
 199     """This is an internal function that is called every time a call is made
 200     during a trace. It keeps track of relationships between calls.
 201     """
 202     global func_count_max
 203     global func_count
 204     global trace_filter
 205     global time_filter
 206     global call_stack
 207     global func_time
 208     global func_time_max
 209
 210     if event == 'call':
 211         keep = True
 212         code = frame.f_code
 213
 214         # Stores all the parts of a human readable name of the current call.
 215         full_name_list = []
 216
 217         # Work out the module name
 218         module = inspect.getmodule(code)
 219         if module:
 220             module_name = module.__name__
 221             module_path = module.__file__
 222             if not settings['include_stdlib'] \
 223                 and is_module_stdlib(module_path):
 224                 keep = False
 225             if module_name == '__main__':
 226                 module_name = ''
 227         else:
 228             module_name = ''
 229         if module_name:
 230             full_name_list.append(module_name)
 231
 232         # Work out the class name.
 233         try:
 234             class_name = frame.f_locals['self'].__class__.__name__
 235             full_name_list.append(class_name)
 236         except (KeyError, AttributeError):
 237             class_name = ''
 238
 239         # Work out the current function or method
 240         func_name = code.co_name
 241         if func_name == '?':
 242             func_name = '__main__'
 243         full_name_list.append(func_name)
 244
 245         # Create a readable representation of the current call
 246         full_name = '.'.join(full_name_list)
 247
 248         # Load the trace filter, if any. 'keep' determines if we should ignore
 249         # this call
 250         if keep and trace_filter:
 251             keep = trace_filter(call_stack, module_name, class_name,
 252                 func_name, full_name)
 253
 254         # Store the call information
 255         if keep:
 256
 257             fr = call_stack[-1]
 258             if fr not in call_dict:
 259                 call_dict[fr] = {}
 260             if full_name not in call_dict[fr]:
 261                 call_dict[fr][full_name] = 0
 262             call_dict[fr][full_name] += 1
 263
 264             if full_name not in func_count:
 265                 func_count[full_name] = 0
 266             func_count[full_name] += 1
 267             if func_count[full_name] > func_count_max:
 268                 func_count_max = func_count[full_name]
 269
 270             call_stack.append(full_name)
 271             call_stack_timer.append(time.time())
 272
 273         else:
 274             call_stack.append('')
 275             call_stack_timer.append(None)
 276
 277     if event == 'return':
 278         if call_stack:
 279             full_name = call_stack.pop(-1)
 280             t = call_stack_timer.pop(-1)
 281             if t and time_filter(stack=call_stack, full_name=full_name):
 282                 if full_name not in func_time:
 283                     func_time[full_name] = 0
 284                 call_time = (time.time() - t)
 285                 func_time[full_name] += call_time
 286                 if func_time[full_name] > func_time_max:
 287                     func_time_max = func_time[full_name]
 288
 289     return tracer
 290
 291
 292 def get_dot(stop=True):
 293     """Returns a string containing a DOT file. Setting stop to True will cause
 294     the trace to stop.
 295     """
 296     global func_time_max
 297
 298     def frac_calculation(func, count):
 299         global func_count_max
 300         global func_time
 301         global func_time_max
 302         calls_frac = float(count) / func_count_max
 303         try:
 304             total_time = func_time[func]
 305         except KeyError:
 306             total_time = 0
 307         if func_time_max:
 308             total_time_frac = float(total_time) / func_time_max
 309         else:
 310             total_time_frac = 0
 311         return calls_frac, total_time_frac, total_time
 312
 313     if stop:
 314         stop_trace()
 315     ret = ['digraph G {', ]
 316     for comp, comp_attr in graph_attributes.items():
 317         ret.append('%s [' % comp)
 318         for attr, val in comp_attr.items():
 319             ret.append('%(attr)s = "%(val)s",' % locals())
 320         ret.append('];')
 321     for func, hits in func_count.items():
 322         calls_frac, total_time_frac, total_time = frac_calculation(func, hits)
 323         col = settings['node_colour'](calls_frac, total_time_frac)
 324         attribs = ['%s="%s"' % a for a in settings['node_attributes'].items()]
 325         node_str = '"%s" [%s];' % (func, ','.join(attribs))
 326         ret.append(node_str % locals())
 327     for fr_key, fr_val in call_dict.items():
 328         if fr_key == '':
 329             continue
 330         for to_key, to_val in fr_val.items():
 331             calls_frac, total_time_frac, totla_time = \
 332                 frac_calculation(to_key, to_val)
 333             col = settings['edge_colour'](calls_frac, total_time_frac)
 334             edge = '[ color = "%s" ]' % col
 335             ret.append('"%s"->"%s" %s' % (fr_key, to_key, edge))
 336     ret.append('}')
 337     ret = '\n'.join(ret)
 338     return ret
 339
 340
 341 def save_dot(filename):
 342     """Generates a DOT file and writes it into filename."""
 343     open(filename, 'w').write(get_dot())
 344
 345
 346 def make_graph(filename, format=None, tool=None, stop=None):
 347     """This has been changed to make_dot_graph."""
 348     raise PyCallGraphException( \
 349         'make_graph is depricated. Please use make_dot_graph')
 350
 351
 352 def make_dot_graph(filename, format='png', tool='dot', stop=True):
 353     """Creates a graph using a Graphviz tool that supports the dot language. It
 354     will output into a file specified by filename with the format specified.
 355     Setting stop to True will stop the current trace.
 356     """
 357     if stop:
 358         stop_trace()
 359
 360     # create a temporary file to be used for the dot data
 361     fd, tempname = tempfile.mkstemp()
 362     f = os.fdopen(fd, 'w')
 363     f.write(get_dot())
 364     f.close()
 365
 366     # normalize filename
 367     regex_user_expand = re.compile('\A~')
 368     if regex_user_expand.match(filename):
 369         filename = os.path.expanduser(filename)
 370     else:
 371         filename = os.path.expandvars(filename)  # expand, just in case
 372
 373     cmd = '%(tool)s -T%(format)s -o%(filename)s %(tempname)s' % locals()
 374     try:
 375         ret = os.system(cmd)
 376         if ret:
 377             raise PyCallGraphException( \
 378                 'The command "%(cmd)s" failed with error ' \
 379                 'code %(ret)i.' % locals())
 380     finally:
 381         os.unlink(tempname)
 382
 383
 384 def simple_memoize(callable_object):
 385     """Simple memoization for functions without keyword arguments.
 386
 387     This is useful for mapping code objects to module in this context.
 388     inspect.getmodule() requires a number of system calls, which may slow down
 389     the tracing considerably. Caching the mapping from code objects (there is
 390     *one* code object for each function, regardless of how many simultaneous
 391     activations records there are).
 392
 393     In this context we can ignore keyword arguments, but a generic memoizer
 394     ought to take care of that as well.
 395     """
 396
 397     cache = dict()
 398     def wrapper(*rest):
 399         if rest not in cache:
 400             cache[rest] = callable_object(*rest)
 401         return cache[rest]
 402
 403     return wrapper
 404
 405
 406 settings = {}
 407 graph_attributes = {}
 408 reset_settings()
 409 reset_trace()
 410 inspect.getmodule = simple_memoize(inspect.getmodule)