"""Heatmap collector for Python profiling with line-level execution heat visualization."""
import base64
import collections
import html
import importlib.resources
import json
import os
import platform
import site
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
from ._css_utils import get_combined_css
from .stack_collector import StackTraceCollector
# ============================================================================
# Data Classes
# ============================================================================
@dataclass
class FileStats:
"""Statistics for a single profiled file."""
filename: str
module_name: str
module_type: str
total_samples: int
total_self_samples: int
num_lines: int
max_samples: int
max_self_samples: int
percentage: float = 0.0
@dataclass
class TreeNode:
"""Node in the hierarchical file tree structure."""
files: List[FileStats] = field(default_factory=list)
samples: int = 0
count: int = 0
children: Dict[str, 'TreeNode'] = field(default_factory=dict)
@dataclass
class ColorGradient:
"""Configuration for heatmap color gradient calculations."""
# Color stops thresholds
stop_1: float = 0.2 # Blue to cyan transition
stop_2: float = 0.4 # Cyan to green transition
stop_3: float = 0.6 # Green to yellow transition
stop_4: float = 0.8 # Yellow to orange transition
stop_5: float = 1.0 # Orange to red transition
# Alpha (opacity) values
alpha_very_cold: float = 0.3
alpha_cold: float = 0.4
alpha_medium: float = 0.5
alpha_warm: float = 0.6
alpha_hot_base: float = 0.7
alpha_hot_range: float = 0.15
# Gradient multiplier
multiplier: int = 5
# Cache for calculated colors
cache: Dict[float, Tuple[int, int, int, float]] = field(default_factory=dict)
# ============================================================================
# Module Path Analysis
# ============================================================================
def get_python_path_info():
"""Get information about Python installation paths for module extraction.
Returns:
dict: Dictionary containing stdlib path, site-packages paths, and sys.path entries.
"""
info = {
'stdlib': None,
'site_packages': [],
'sys_path': []
}
# Get standard library path from os module location
try:
if hasattr(os, '__file__') and os.__file__:
info['stdlib'] = Path(os.__file__).parent
except (AttributeError, OSError):
pass # Silently continue if we can't determine stdlib path
# Get site-packages directories
site_packages = []
try:
site_packages.extend(Path(p) for p in site.getsitepackages())
except (AttributeError, OSError):
pass # Continue without site packages if unavailable
# Get user site-packages
try:
user_site = site.getusersitepackages()
if user_site and Path(user_site).exists():
site_packages.append(Path(user_site))
except (AttributeError, OSError):
pass # Continue without user site packages
info['site_packages'] = site_packages
info['sys_path'] = [Path(p) for p in sys.path if p]
return info
def extract_module_name(filename, path_info):
"""Extract Python module name and type from file path.
Args:
filename: Path to the Python file
path_info: Dictionary from get_python_path_info()
Returns:
tuple: (module_name, module_type) where module_type is one of:
'stdlib', 'site-packages', 'project', or 'other'
"""
if not filename:
return ('unknown', 'other')
try:
file_path = Path(filename)
except (ValueError, OSError):
return (str(filename), 'other')
# Check if it's in stdlib
if path_info['stdlib'] and _is_subpath(file_path, path_info['stdlib']):
try:
rel_path = file_path.relative_to(path_info['stdlib'])
return (_path_to_module(rel_path), 'stdlib')
except ValueError:
pass
# Check site-packages
for site_pkg in path_info['site_packages']:
if _is_subpath(file_path, site_pkg):
try:
rel_path = file_path.relative_to(site_pkg)
return (_path_to_module(rel_path), 'site-packages')
except ValueError:
continue
# Check other sys.path entries (project files)
if not str(file_path).startswith(('<', '[')): # Skip special files
for path_entry in path_info['sys_path']:
if _is_subpath(file_path, path_entry):
try:
rel_path = file_path.relative_to(path_entry)
return (_path_to_module(rel_path), 'project')
except ValueError:
continue
# Fallback: just use the filename
return (_path_to_module(file_path), 'other')
def _is_subpath(file_path, parent_path):
try:
file_path.relative_to(parent_path)
return True
except (ValueError, OSError):
return False
def _path_to_module(path):
if isinstance(path, str):
path = Path(path)
# Remove .py extension
if path.suffix == '.py':
path = path.with_suffix('')
# Convert path separators to dots
parts = path.parts
# Handle __init__ files - they represent the package itself
if parts and parts[-1] == '__init__':
parts = parts[:-1]
return '.'.join(parts) if parts else path.stem
# ============================================================================
# Helper Classes
# ============================================================================
class _TemplateLoader:
"""Loads and caches HTML/CSS/JS templates for heatmap generation."""
def __init__(self):
"""Load all templates and assets once."""
self.index_template = None
self.file_template = None
self.index_css = None
self.index_js = None
self.file_css = None
self.file_js = None
self.logo_html = None
self._load_templates()
def _load_templates(self):
"""Load all template files from _heatmap_assets."""
try:
template_dir = importlib.resources.files(__package__)
assets_dir = template_dir / "_heatmap_assets"
# Load HTML templates
self.index_template = (assets_dir / "heatmap_index_template.html").read_text(encoding="utf-8")
self.file_template = (assets_dir / "heatmap_pyfile_template.html").read_text(encoding="utf-8")
# Load CSS (same file used for both index and file pages)
css_content = get_combined_css("heatmap")
self.index_css = css_content
self.file_css = css_content
# Load JS
self.index_js = (assets_dir / "heatmap_index.js").read_text(encoding="utf-8")
self.file_js = (assets_dir / "heatmap.js").read_text(encoding="utf-8")
# Load Python logo
logo_dir = template_dir / "_assets"
try:
png_path = logo_dir / "python-logo-only.png"
b64_logo = base64.b64encode(png_path.read_bytes()).decode("ascii")
self.logo_html = f''
except (FileNotFoundError, IOError) as e:
self.logo_html = '
'
print(f"Warning: Could not load Python logo: {e}")
except (FileNotFoundError, IOError) as e:
raise RuntimeError(f"Failed to load heatmap template files: {e}") from e
class _TreeBuilder:
"""Builds hierarchical tree structure from file statistics."""
@staticmethod
def build_file_tree(file_stats: List[FileStats]) -> Dict[str, TreeNode]:
"""Build hierarchical tree grouped by module type, then by module structure.
Args:
file_stats: List of FileStats objects
Returns:
Dictionary mapping module types to their tree roots
"""
# Group by module type first
type_groups = {'stdlib': [], 'site-packages': [], 'project': [], 'other': []}
for stat in file_stats:
type_groups[stat.module_type].append(stat)
# Build tree for each type
trees = {}
for module_type, stats in type_groups.items():
if not stats:
continue
root_node = TreeNode()
for stat in stats:
module_name = stat.module_name
parts = module_name.split('.')
# Navigate/create tree structure
current_node = root_node
for i, part in enumerate(parts):
if i == len(parts) - 1:
# Last part - store the file
current_node.files.append(stat)
else:
# Intermediate part - create or navigate
if part not in current_node.children:
current_node.children[part] = TreeNode()
current_node = current_node.children[part]
# Calculate aggregate stats for this type's tree
_TreeBuilder._calculate_node_stats(root_node)
trees[module_type] = root_node
return trees
@staticmethod
def _calculate_node_stats(node: TreeNode) -> Tuple[int, int]:
"""Recursively calculate aggregate statistics for tree nodes.
Args:
node: TreeNode to calculate stats for
Returns:
Tuple of (total_samples, file_count)
"""
total_samples = 0
file_count = 0
# Count files at this level
for file_stat in node.files:
total_samples += file_stat.total_samples
file_count += 1
# Recursively process children
for child in node.children.values():
child_samples, child_count = _TreeBuilder._calculate_node_stats(child)
total_samples += child_samples
file_count += child_count
node.samples = total_samples
node.count = file_count
return total_samples, file_count
class _HtmlRenderer:
"""Renders hierarchical tree structures as HTML."""
def __init__(self, file_index: Dict[str, str], color_gradient: ColorGradient,
calculate_intensity_color_func):
"""Initialize renderer with file index and color calculation function.
Args:
file_index: Mapping from filenames to HTML file names
color_gradient: ColorGradient configuration
calculate_intensity_color_func: Function to calculate colors
"""
self.file_index = file_index
self.color_gradient = color_gradient
self.calculate_intensity_color = calculate_intensity_color_func
self.heatmap_bar_height = 16
def render_hierarchical_html(self, trees: Dict[str, TreeNode]) -> str:
"""Build hierarchical HTML with type sections and collapsible module folders.
Args:
trees: Dictionary mapping module types to tree roots
Returns:
Complete HTML string for all sections
"""
type_names = {
'stdlib': '📚 Standard Library',
'site-packages': '📦 Site Packages',
'project': '🏗️ Project Files',
'other': '📄 Other Files'
}
sections = []
for module_type in ['project', 'stdlib', 'site-packages', 'other']:
if module_type not in trees:
continue
tree = trees[module_type]
# Project starts expanded, others start collapsed
is_collapsed = module_type in {'stdlib', 'site-packages', 'other'}
icon = '▶' if is_collapsed else '▼'
content_style = ' style="display: none;"' if is_collapsed else ''
section_html = f'''
'''
# Render root folders
root_folders = sorted(tree.children.items(),
key=lambda x: x[1].samples, reverse=True)
for folder_name, folder_node in root_folders:
section_html += self._render_folder(folder_node, folder_name, level=1)
# Render root files (files not in any module)
if tree.files:
sorted_files = sorted(tree.files, key=lambda x: x.total_samples, reverse=True)
section_html += '
\n'
for stat in sorted_files:
section_html += self._render_file_item(stat, indent=' ')
section_html += '
\n'
section_html += '
\n
\n'
sections.append(section_html)
return '\n'.join(sections)
def _render_folder(self, node: TreeNode, name: str, level: int = 1) -> str:
"""Render a single folder node recursively.
Args:
node: TreeNode to render
name: Display name for the folder
level: Nesting level for indentation
Returns:
HTML string for this folder and its contents
"""
indent = ' ' * level
parts = []
# Render folder header (collapsed by default)
parts.append(f'{indent}