"""Heatmap collector for Python profiling with line-level execution heat visualization."""
import base64
import collections
import html
import importlib.resources
import json
import math
import os
import platform
import site
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Tuple
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)
# ============================================================================
# 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
shared_js = (assets_dir / "heatmap_shared.js").read_text(encoding="utf-8")
self.index_js = f"{shared_js}\n{(assets_dir / 'heatmap_index.js').read_text(encoding='utf-8')}"
self.file_js = f"{shared_js}\n{(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 = '