generate gprof-like stats (work in progress)

prof_wrangler.py

#!/usr/bin/env python3

import os
import json
import collections
import multiprocessing
import typing

dirp = '/tmp/prof'


def iter_files(dirp):
    for fn in os.listdir(dirp):
        if fn.endswith('.json'):
            yield os.path.join(dirp, fn)


def parse_name(name):
    tokens = name.split(':')
    if len(tokens) >= 3:
        return tokens[2]
    return name


def traverse(d):
    frames = d['shared']['frames']
    for event in d['profiles'][0]['events']:
        frame = frames[event['frame']]
        if event['type'] == 'C':
            frame['close'] = event['at']
        elif event['type'] == 'O':
            frame['open'] = event['at']
    for f in frames:
        yield parse_name(f['name']), f['close'] - f['open']


class CallCount(typing.NamedTuple):
    count: int
    cost: int

    @staticmethod
    def empty():
        return CallCount(0, 0)


def process_one(fn):
    rec = collections.defaultdict(CallCount.empty)
    with open(fn) as fp:
        d = json.load(fp)
        for func, cost in traverse(d):
            (count, cost_) = rec[func]
            rec[func] = CallCount(count + 1, cost_ + cost)
    return rec


def show_percentage(v):
    if v < 0.000001:
        return '0.0%'
    return '{:.2}%'.format(100 * v)


def show_float(v):
    return '{:.2}'.format(v)


def display(items, sort_func=None, limit=25):
    total_cost = sum(call_count[1] for _, call_count in items)
    if sort_func is None:
        sorted_items = sorted(items, key=lambda cc: cc[1], reverse=True)
    else:
        sorted_items = sort_func(items)
    longest = max(len(func) for func, _ in sorted_items[:limit])
    print('name'.rjust(longest + 4), 'calls'.rjust(8), 'cumula'.rjust(12),
          'per call'.rjust(8), '    %')
    for func, cc in sorted_items[:limit]:
        print(func.rjust(longest + 4),
              str(cc[0]).rjust(8),
              str(cc[1]).rjust(12),
              show_float(cc[1] / cc[0]).rjust(8),
              show_percentage(cc[1] / total_cost).rjust(8))


def main(dirp):
    rec = collections.defaultdict(CallCount.empty)
    pool = multiprocessing.Pool()
    it = iter_files(dirp)
    for each_rec in pool.map(process_one, it):
        for func, call_count in each_rec.items():
            call_count_ = rec[func]
            rec[func] = CallCount(call_count[0] + call_count_[0],
                                  call_count[1] + call_count_[1])
    display(rec.items())


if __name__ == '__main__':
    main(dirp)