Merge pull request gak#124 from sMAshdot/develop

Specify Custom Groups

Author: gak

examples/graphviz/example_with_submodules/__init__.py

@@ -0,0 +1,3 @@
+from .example_with_submodules import main
+
+__all__ = [main]

examples/graphviz/example_with_submodules/example_with_submodules.py

@@ -0,0 +1,13 @@
+from submodule_one import SubmoduleOne
+from submodule_two import SubmoduleTwo
+
+
+def main():
+    s1 = SubmoduleOne()
+    s1.report()
+
+    s2 = SubmoduleTwo()
+    s2.report()
+
+if __name__ == "__main__":
+    main()

examples/graphviz/example_with_submodules/helpers.py

@@ -0,0 +1,2 @@
+def helper(something):
+    return something

examples/graphviz/example_with_submodules/submodule_one.py

@@ -0,0 +1,6 @@
+class SubmoduleOne(object):
+    def __init__(self):
+        self.one = 1
+
+    def report(self):
+        return self.one

examples/graphviz/example_with_submodules/submodule_two.py

@@ -0,0 +1,9 @@
+from helpers import helper
+
+
+class SubmoduleTwo(object):
+    def __init__(self):
+        self.two = 2
+
+    def report(self):
+        return helper(self.two)

examples/graphviz/grouper.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+'''
+This example demonstrates the use of grouping.
+'''
+from pycallgraph import PyCallGraph
+from pycallgraph import Config
+from pycallgraph import GlobbingFilter
+from pycallgraph import Grouper
+from pycallgraph.output import GraphvizOutput
+import example_with_submodules
+
+
+def run(name, trace_grouper=None, config=None, comment=None):
+    if not config:
+        config = Config()
+
+    config.trace_filter = GlobbingFilter()
+
+    if trace_grouper is not None:
+        config.trace_grouper = trace_grouper
+
+    graphviz = GraphvizOutput()
+    graphviz.output_file = 'grouper-{}.png'.format(name)
+    if comment:
+        graphviz.graph_attributes['graph']['label'] = comment
+
+    with PyCallGraph(config=config, output=graphviz):
+        example_with_submodules.main()
+
+
+def group_none():
+    run(
+        'without',
+        comment='Default grouping.'
+    )
+
+
+def group_some():
+    trace_grouper = Grouper(groups=[
+        'example_with_submodules.submodule_one.*',
+        'example_with_submodules.submodule_two.*',
+        'example_with_submodules.helpers.*',
+    ])
+
+    run(
+        'with',
+        trace_grouper=trace_grouper,
+        comment='Should assign groups to the two submodules.',
+    )
+
+
+def group_methods():
+    trace_grouper = Grouper(groups=[
+        'example_with_submodules.*.report',
+        ])
+
+    run(
+        'methods',
+        trace_grouper=trace_grouper,
+        comment='Should assign a group to the report methods.',
+    )
+
+
+def main():
+    group_none()
+    group_some()
+    group_methods()
+
+
+if __name__ == '__main__':
+    main()

pycallgraph/__init__.py

@@ -16,6 +16,7 @@
 from .exceptions import PyCallGraphException
 from .config import Config
 from .globbing_filter import GlobbingFilter
+from .grouper import Grouper
 from .util import Util
 from .color import Color
 from .color import ColorException

pycallgraph/config.py

@@ -3,6 +3,7 @@
 
 from .output import outputters
 from .globbing_filter import GlobbingFilter
+from .grouper import Grouper
 
 
 class Config(object):
@@ -31,6 +32,9 @@ def __init__(self, **kwargs):
             include=['*'],
         )
 
+        # Grouping
+        self.trace_grouper = Grouper()
+
         self.did_init = True
 
         # Update the defaults with anything from kwargs

pycallgraph/grouper.py

@@ -0,0 +1,26 @@
+from fnmatch import fnmatch
+
+
+class Grouper(object):
+    '''Group module names.
+
+    By default, objects are grouped by their top-level module name. Additional
+    groups can be specified with the groups list and all objects will be
+    matched against it.
+    '''
+
+    def __init__(self, groups=None):
+        if groups is None:
+            groups = []
+
+        self.groups = groups
+
+    def __call__(self, full_name=None):
+        for pattern in self.groups:
+            if fnmatch(full_name, pattern):
+                if pattern[-2:] == ".*":
+                    # a wilcard in the middle is probably meaningful, while at
+                    # the end, it's only noise and can be removed
+                    return pattern[:-2]
+                return pattern
+        return full_name.split('.', 1)[0]

pycallgraph/tracer.py

@@ -290,20 +290,17 @@ def __getstate__(self):
 
         return odict
 
-    def group(self, name):
-        return name.split('.', 1)[0]
-
     def groups(self):
         grp = defaultdict(list)
         for node in self.nodes():
-            grp[self.group(node.name)].append(node)
+            grp[node.group].append(node)
         for g in grp.iteritems():
             yield g
 
     def stat_group_from_func(self, func, calls):
         stat_group = StatGroup()
         stat_group.name = func
-        stat_group.group = self.group(func)
+        stat_group.group = self.config.trace_grouper(func)
         stat_group.calls = Stat(calls, self.func_count_max)
         stat_group.time = Stat(self.func_time.get(func, 0), self.func_time_max)
         stat_group.memory_in = Stat(