# Copyright 2017-2023 Lawrence Livermore National Security, LLC and other
# Hatchet Project Developers. See the top-level LICENSE file for details.
#
# SPDX-License-Identifier: MIT
import json
import pandas as pd
import hatchet.graphframe
from hatchet.node import Node
from hatchet.graph import Graph
from hatchet.frame import Frame
[docs]class PyinstrumentReader:
def __init__(self, filename):
self.pyinstrument_json_filename = filename
self.graph_dict = {}
self.list_roots = []
self.node_dicts = []
[docs] def create_graph(self):
def parse_node_literal(child_dict, hparent):
"""Create node_dict for one node and then call the function
recursively on all children."""
hnode = Node(
Frame({"name": child_dict["function"], "type": "function"}), hparent
)
child_node_dict = {
"node": hnode,
"name": child_dict["function"],
"file": child_dict["file_path_short"],
"line": child_dict["line_no"],
"time": child_dict["time"],
"time (inc)": child_dict["time"],
"is_application_code": child_dict["is_application_code"],
}
hparent.add_child(hnode)
self.node_dicts.append(child_node_dict)
if "children" in child_dict:
for child in child_dict["children"]:
# Pyinstrument's time metric actually stores inclusive time.
# To calculate exclusive time, we subtract the children's time
# from the parent's time.
child_node_dict["time"] -= child["time"]
parse_node_literal(child, hnode)
# start with creating a node_dict for each root
graph_root = Node(
Frame(
{"name": self.graph_dict["root_frame"]["function"], "type": "function"}
),
None,
)
node_dict = {
"node": graph_root,
"name": self.graph_dict["root_frame"]["function"],
"file": self.graph_dict["root_frame"]["file_path_short"],
"line": self.graph_dict["root_frame"]["line_no"],
"time": self.graph_dict["root_frame"]["time"],
"time (inc)": self.graph_dict["root_frame"]["time"],
"is_application_code": self.graph_dict["root_frame"]["is_application_code"],
}
self.node_dicts.append(node_dict)
self.list_roots.append(graph_root)
# call recursively on all children of root
if "children" in self.graph_dict["root_frame"]:
for child in self.graph_dict["root_frame"]["children"]:
# Pyinstrument's time metric actually stores inclusive time.
# To calculate exclusive time, we subtract the children's time
# from the parent's time.
node_dict["time"] -= child["time"]
parse_node_literal(child, graph_root)
graph = Graph(self.list_roots)
graph.enumerate_traverse()
return graph
[docs] def read(self):
with open(self.pyinstrument_json_filename) as pyinstrument_json:
self.graph_dict = json.load(pyinstrument_json)
graph = self.create_graph()
dataframe = pd.DataFrame(data=self.node_dicts)
dataframe.set_index(["node"], inplace=True)
dataframe.sort_index(inplace=True)
return hatchet.graphframe.GraphFrame(graph, dataframe, ["time"], ["time (inc)"])