Convert github.com/joerick/pyinstrument's json format to flamegraph

pyinstrument_json_to_flamegraph.py

#!/usr/bin/env python3
# Author: Xie Yanbo <xieyanbo@gmail.com>

import json

USE_LONG_FILE_PATH = True

def to_flamegraph(pyinst_data):

    def walk(data):
        if USE_LONG_FILE_PATH:
            file_path = data['file_path']
        else:
            file_path = data['file_path_short']
        id = '{}:{}:{}'.format(file_path,
                               data['line_no'],
                               data['function'])
        time = data['time']
        children_time = 0
        if data['children']:
            for child in data['children']:
                children_time += child['time']
                for cid, ctime in walk(child):
                    yield '{};{}'.format(id, cid), ctime

        yield id, time - children_time

    root = pyinst_data['root_frame']
    result = []
    for id, time in walk(root):
        line = '{} {}'.format(id, int(time * 1000))
        result.append(line)

    return '\n'.join(result) + '\n'


def help():
    print("Convert github.com/joerick/pyinstrument's json format to flamegraph")
    print('Usage: {} [-s] <input_file> <output_file>'.format(sys.argv[0]))
    print('  -s, --short-file-path  Use short file path (defulat: false)')


if __name__ == '__main__':
    import sys
    if '-h' in sys.argv or '--help' in sys.argv:
        help()
        sys.exit(0)

    if '-s' in sys.argv or '--short-file-path' in sys.argv:
        USE_LONG_FILE_PATH = False
        try:
            sys.argv.remove('-s')
        except ValueError:
            pass
        try:
            sys.argv.remove('--short-file-path')
        except ValueError:
            pass

    if len(sys.argv) >= 2:
        input_file = open(sys.argv[1])
    else:
        input_file = sys.stdin
    if len(sys.argv) >= 3:
        output_file = open(sys.argv[2], 'w')
    else:
        output_file = sys.stdout

    output_file.write(to_flamegraph(json.load(input_file)))