#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
#

'''
calltree.py - Call tree manipulator

This script manipulates call tree generated by JASMIN CallTreeManager. It
prints, folds, and limit depth of the input call tree.
'''

import argparse
import copy
import json

from collections import OrderedDict


def fold_tree(tree):

    def fold_tree_recursive(tree, compare_equal):
        if not tree:
            return None
        new_tree = OrderedDict()
        new_tree["id"] = tree["id"]
        new_tree["cycle"] = tree["cycle"]
        new_tree["children"] = [fold_tree_recursive(t, compare_equal) for t in
                                tree["children"]]
        if not new_tree["children"]:
            return new_tree
        new_children = [new_tree["children"][0]]
        for i in xrange(1, len(new_tree["children"])):
            if not compare_equal(new_tree["children"][i], new_children[-1]):
                new_children.append(new_tree["children"][i])
        new_tree["children"] = new_children
        return new_tree

    def tree_equal(src, dst):
        if src["id"] != dst["id"]:
            return False
        elif len(src["children"]) != len(dst["children"]):
            return False
        else:
            for x, y in zip(src["children"], dst["children"]):
                if not tree_equal(x, y):
                    return False
            return True

    return fold_tree_recursive(tree, tree_equal)


def print_tree(tree, max_level=None):
    def print_tree_recursive(tree, level, max_level):
        if max_level and level > max_level:
            return
        print "--" * level, tree["id"], tree["cycle"]
        for c in tree["children"]:
            print_tree_recursive(c, level + 1, max_level)
    print_tree_recursive(tree, 1, max_level)


def main():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("call_tree",
                        help="Call tree file (json) generated by JASMIN")
    parser.add_argument("--save", help="Output filename", default=None)
    parser.add_argument("--max-depth", default=None, type=int,
                        help="Max depth for output tree")
    parser.add_argument("--fold", action="store_true",
                        help="Try hard to fold tree")

    args = parser.parse_args()

    tree = json.load(file(args.call_tree), object_hook=OrderedDict)
    if tree["id"] == "ROOT":
        tree = tree["children"][0]
    if args.fold:
        tree = fold_tree(tree)
    if args.save:
        json.dump(tree, file(args.save, "w"), indent=2)
    print_tree(tree, args.max_depth)


if __name__ == "__main__":
    main()

Add bentoo-calltree tool