xml-timings-to-graph

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright © 2019 Martin Ueding <dev@martin-ueding.de>

import argparse
import collections
import hashlib
import itertools
import math
import os
import pprint
import subprocess
import sys

import jinja2
import lxml.etree
import tqdm


Vertex = collections.namedtuple('Vertex', ['id', 'function', 'info', 'cumtime', 'selftime', 'calls', 'hue'])
Edge = collections.namedtuple('Edge', ['from_id', 'to_id', 'function', 'info', 'cumtime', 'calls'])


def node_id(function, info):
    return 'n{}'.format(hashlib.sha1((function + ' ' + info).encode()).hexdigest())


def get_info(call):
    info = call.attrib['info']

    parent = call
    while True:
        parent = parent.find('..')
        if parent is None:
            break
        parent_info = parent.get('info', '')
        if parent_info != info:
            infos = [parent_info, info]
            info = '/'.join(i for i in infos if len(i) > 0)

    return info


def get_edge(child, caller_id):
    info = get_info(child)

    return Edge(
        caller_id,
        node_id(child.attrib['function'], info),
        child.attrib['function'],
        info,
        float(child.find('./total').attrib['time']),
        1)


def parse_call(call):
    cumtime = float(call.find('./total').attrib['time'])
    function = call.attrib['function']
    info = get_info(call)
    id = node_id(function, info)

    child_calls = call.findall('./call')
    edges = [
        get_edge(child, id)
        for child in child_calls
    ]

    selftime = cumtime - sum(edge.cumtime for edge in edges)

    vertex = Vertex(id=id, function=function, info=info, cumtime=cumtime, selftime=selftime, calls=1, hue=None)

    return vertex, edges


def main():
    options = _parse_args()

    env = jinja2.Environment(loader=jinja2.FileSystemLoader(
        os.path.dirname(os.path.abspath(__file__))))
    template = env.get_template('profile.dot.j2')

    tree = lxml.etree.parse(options.profile)
    print('Loaded XML file.')

    calls = tree.findall('.//call')
    print('Found {} calls.'.format(len(calls)))

    tottime = float(tree.find('.//call[@function="main"]/total').attrib['time'])

    vertices = []
    edges = []

    for call in tqdm.tqdm(calls):
        v, e = parse_call(call)
        vertices.append(v)
        edges += e

    print('Gathered {} vertices and {} edges.'.format(len(vertices), len(edges)))

    vertex_keyfun = lambda v: v.id
    vertices.sort(key=vertex_keyfun)
    grouped_vertices = []
    for vertex_id, group in itertools.groupby(vertices, vertex_keyfun):
        group = list(group)
        v = group[0]

        cumtime = sum(g.cumtime for g in group)
        selftime = sum(g.selftime for g in group)

        grouped_vertices.append(Vertex(
            id=v.id, function=v.function, info=v.info,
            cumtime=cumtime,
            selftime=selftime,
            calls=len(group),
            hue='{:.4f}'.format(250/360 * (1 - math.pow(cumtime / tottime, 0.5))),
        ))

    print('Grouped vertices down to {}.'.format(len(grouped_vertices)))

    edge_keyfun = lambda e: (e.from_id, e.to_id)
    edges.sort(key=edge_keyfun)
    grouped_edges = []
    for (from_id, to_id), group in itertools.groupby(edges, edge_keyfun):
        group = list(group)
        e = group[0]

        grouped_edges.append(Edge(
            from_id=e.from_id, to_id=e.to_id,
            function=e.function, info=e.info,
            cumtime=sum(g.cumtime for g in group),
            calls=len(group)))

    print('Grouped edges down to {}.'.format(len(grouped_edges)))

    rendered = template.render(vertices=grouped_vertices, edges=grouped_edges, tottime=tottime)
    print('Generated graph.')

    with open('profile.dot', 'w') as f:
        f.write(rendered)
    print('Wrote graph.')

    subprocess.run(['dot', 'profile.dot', '-Tpdf', '-o', 'profile.pdf'])
    subprocess.run(['dot', 'profile.dot', '-Tpng', '-o', 'profile.png'])
    print('Compiled graph.')



def _parse_args():
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('profile')
    options = parser.parse_args()

    return options


if __name__ == '__main__':
    main()