130 lines
4.2 KiB
Python
130 lines
4.2 KiB
Python
from __future__ import annotations
|
|
import time
|
|
from typing import Any
|
|
from contextlib import contextmanager
|
|
|
|
def get_node_type(node: Any) -> str:
|
|
if hasattr(node, "bl_label") and node.bl_label:
|
|
return str(node.bl_label)
|
|
if hasattr(node, "bl_idname") and node.bl_idname:
|
|
return str(node.bl_idname)
|
|
return type(node).__name__
|
|
|
|
class NodeProfiler:
|
|
def __init__(self) -> None:
|
|
# Maps node_type -> [execution_count, total_self_time_seconds]
|
|
self.stats: dict[str, list[float]] = {}
|
|
self._stack: list[dict[str, Any]] = []
|
|
|
|
@contextmanager
|
|
def profile_node(self, node: Any):
|
|
node_type = get_node_type(node)
|
|
now = time.perf_counter()
|
|
|
|
# Pause parent node's self-timer
|
|
if self._stack:
|
|
parent = self._stack[-1]
|
|
parent["self_time"] += (now - parent["start_time"])
|
|
|
|
frame = {
|
|
"node_type": node_type,
|
|
"start_time": now,
|
|
"self_time": 0.0,
|
|
}
|
|
self._stack.append(frame)
|
|
try:
|
|
yield
|
|
finally:
|
|
end_time = time.perf_counter()
|
|
popped = self._stack.pop()
|
|
total_node_self_time = popped["self_time"] + (end_time - popped["start_time"])
|
|
|
|
if node_type not in self.stats:
|
|
self.stats[node_type] = [0.0, 0.0]
|
|
self.stats[node_type][0] += 1.0
|
|
self.stats[node_type][1] += total_node_self_time
|
|
|
|
# Resume parent node's timer
|
|
if self._stack:
|
|
self._stack[-1]["start_time"] = time.perf_counter()
|
|
|
|
def format_table(self) -> str:
|
|
headers = ["node_type", "execution_count", "avg_execution_time", "tot_execution_time"]
|
|
|
|
if not self.stats:
|
|
col_widths = [20, 15, 18, 18]
|
|
header_line = " | ".join(headers[i].ljust(col_widths[i]) for i in range(4))
|
|
separator_line = "-+-".join("-" * col_widths[i] for i in range(4))
|
|
return f"{header_line}\n{separator_line}\n(No nodes executed)"
|
|
|
|
sorted_stats = sorted(self.stats.items(), key=lambda item: item[1][1], reverse=True)
|
|
|
|
rows: list[list[str]] = []
|
|
tot_count = 0
|
|
tot_time_all = 0.0
|
|
|
|
for node_type, (count, tot_time) in sorted_stats:
|
|
icount = int(count)
|
|
tot_count += icount
|
|
tot_time_all += tot_time
|
|
avg_time = tot_time / icount if icount > 0 else 0.0
|
|
rows.append([
|
|
node_type,
|
|
str(icount),
|
|
f"{avg_time:.6f}s",
|
|
f"{tot_time:.6f}s",
|
|
])
|
|
|
|
col_widths = [
|
|
max(len(headers[i]), max(len(row[i]) for row in rows))
|
|
for i in range(4)
|
|
]
|
|
col_widths[0] = max(col_widths[0], len("Total"))
|
|
col_widths[1] = max(col_widths[1], len(str(tot_count)))
|
|
col_widths[3] = max(col_widths[3], len(f"{tot_time_all:.6f}s"))
|
|
|
|
lines: list[str] = []
|
|
header_line = " | ".join(headers[i].ljust(col_widths[i]) for i in range(4))
|
|
separator_line = "-+-".join("-" * col_widths[i] for i in range(4))
|
|
|
|
lines.append(header_line)
|
|
lines.append(separator_line)
|
|
for row in rows:
|
|
lines.append(" | ".join([
|
|
row[0].ljust(col_widths[0]),
|
|
row[1].rjust(col_widths[1]),
|
|
row[2].rjust(col_widths[2]),
|
|
row[3].rjust(col_widths[3]),
|
|
]))
|
|
lines.append(separator_line)
|
|
total_row = [
|
|
"Total".ljust(col_widths[0]),
|
|
str(tot_count).rjust(col_widths[1]),
|
|
"-".rjust(col_widths[2]),
|
|
f"{tot_time_all:.6f}s".rjust(col_widths[3]),
|
|
]
|
|
lines.append(" | ".join(total_row))
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
_current_profiler: NodeProfiler | None = None
|
|
|
|
@contextmanager
|
|
def profiler_context():
|
|
global _current_profiler
|
|
old = _current_profiler
|
|
profiler = NodeProfiler()
|
|
_current_profiler = profiler
|
|
try:
|
|
yield profiler
|
|
finally:
|
|
_current_profiler = old
|
|
|
|
@contextmanager
|
|
def profile_node_execution(node: Any):
|
|
if _current_profiler is None:
|
|
yield
|
|
return
|
|
with _current_profiler.profile_node(node):
|
|
yield
|