# -*- coding: utf-8 -*-
"""
pyheat.pyheat
~~~~~~~~~~~
Class definitions to view python file as a heatmap highlighting
the most time consuming parts of the file.
"""
import matplotlib.pyplot as plt
import numpy as np
import pprofile
from matplotlib.widgets import Slider
class PyHeat(object):
"""Class to display the python file provided as a heatmap highlighting the
most time consuming parts of the file.
"""
class FileDetails(object):
"""Class to keep track of useful file related details used in creating
heatmaps.
"""
def __init__(self, filepath):
"""Constructor.
@param filepath: Path of the file to profile.
"""
self.path = filepath
self.lines = []
self.length = 0
self.data = None
def __init__(self, filepath):
"""Constructor.
@param filepath: Path of the file to profile.
"""
self.pyfile = PyHeat.FileDetails(filepath)
self.line_profiler = None
def create_heatmap(self):
"""Method to define the heatmap using profile data."""
# Profile the file.
self.__profile_file()
# Map profile stats to heatmap data.
self.__fetch_heatmap_data_from_profile()
# Create heatmap.
self.__create_heatmap_plot()
def show_heatmap(self, blocking=True, output_file=None, enable_scroll=False):
"""Method to actually display the heatmap created.
@param blocking: When set to False makes an unblocking plot show.
@param output_file: If not None the heatmap image is output to this
file. Supported formats: (eps, pdf, pgf, png, ps, raw, rgba, svg,
svgz)
@param enable_scroll: Flag used add a scroll bar to scroll long files.
"""
if output_file is None:
if enable_scroll:
# Add a new axes which will be used as scroll bar.
axpos = plt.axes([0.12, 0.1, 0.625, 0.03])
spos = Slider(axpos, "Scroll", 10, len(self.pyfile.lines))
def update(val):
"""Method to update position when slider is moved."""
pos = spos.val
self.ax.axis([0, 1, pos, pos - 10])
self.fig.canvas.draw_idle()
spos.on_changed(update)
plt.show(block=blocking)
else:
plt.savefig(output_file)
def close_heatmap(self):
"""Method to close the heatmap display created."""
plt.close("all")
def __profile_file(self):
"""Method used to profile the given file line by line."""
self.line_profiler = pprofile.Profile()
self.line_profiler.runfile(
open(self.pyfile.path, "r"), {}, self.pyfile.path
)
def __get_line_profile_data(self):
"""Method to procure line profiles.
@return: Line profiles if the file has been profiles else empty
dictionary.
"""
if self.line_profiler is None:
return {}
# the [0] is because pprofile.Profile.file_dict stores the line_dict
# in a list so that it can be modified in a thread-safe way
# see https://github.com/vpelletier/pprofile/blob/da3d60a1b59a061a0e2113bf768b7cb4bf002ccb/pprofile.py#L398
return self.line_profiler.file_dict[self.pyfile.path][0].line_dict
def __fetch_heatmap_data_from_profile(self):
"""Method to create heatmap data from profile information."""
# Read lines from file.
with open(self.pyfile.path, "r") as file_to_read:
for line in file_to_read:
# Remove return char from the end of the line and add a
# space in the beginning for better visibility.
self.pyfile.lines.append(" " + line.strip("\n"))
# Total number of lines in file.
self.pyfile.length = len(self.pyfile.lines)
# Fetch line profiles.
line_profiles = self.__get_line_profile_data()
# Creating an array of data points. As the profile keys are 1 indexed
# we should range from 1 to line_count + 1 and not 0 to line_count.
arr = []
for line_num in range(1, self.pyfile.length + 1):
if line_num in line_profiles:
# line_profiles[i] will have multiple entries if line i is
# invoked from multiple places in the code. Here we sum over
# each invocation to get the total time spent on that line.
line_times = [
ltime for _, ltime in line_profiles[line_num].values()
]
arr.append([sum(line_times)])
else:
arr.append([0.0])
# Create nd-array from list of data points.
self.pyfile.data = np.array(arr)
def __create_heatmap_plot(self):
"""Method to actually create the heatmap from profile stats."""
# Define the heatmap plot.
height = len(self.pyfile.lines) / 3
width = max(map(lambda x: len(x), self.pyfile.lines)) / 8
self.fig, self.ax = plt.subplots(figsize=(width, height))
# Set second sub plot to occupy bottom 20%
plt.subplots_adjust(bottom=0.20)
# Heat scale orange to red
heatmap = self.ax.pcolor(self.pyfile.data, cmap="OrRd")
# X Axis
# Remove X axis.
self.ax.xaxis.set_visible(False)
# Y Axis
# Create lables for y-axis ticks
row_labels = range(1, self.pyfile.length + 1)
# Set y-tick labels.
self.ax.set_yticklabels(row_labels, minor=False)
# Put y-axis major ticks at the middle of each cell.
self.ax.set_yticks(np.arange(self.pyfile.data.shape[0]) + 0.5, minor=False)
# Inver y-axis to have top down line numbers
self.ax.invert_yaxis()
# Plot definitions
# Set plot y-axis label.
plt.ylabel("Line Number")
# Annotate each cell with lines in file in order.
max_time_spent_on_a_line = max(self.pyfile.data)
for i, line in enumerate(self.pyfile.lines):
# In order to ensure easy readability of the code, we need to
# invert colour of text display for darker colours which
# correspond to higher amount of time spent on the line.
if self.pyfile.data[i] >= 0.7 * max_time_spent_on_a_line:
color = (1.0, 1.0, 1.0) # White text
else:
color = (0.0, 0.0, 0.0) # Black text
plt.text(
0.0,
i + 0.5,
line,
ha="left",
va="center",
color=color,
clip_on=True,
)
# Define legend
cbar = plt.colorbar(heatmap)
cbar.set_label("# of seconds")
