Files
2125_GCE/xic_extensions/profiler.py
T
2026-07-09 12:54:44 -04:00

56 lines
1.7 KiB
Python
Executable File

import time
from typing import Dict, Optional
from dataclasses import dataclass
@dataclass
class ProfileSnapshot:
elapsed: float
call_count: int
memory_approx: int = 0
class SegmentProfiler:
def __init__(self):
self._profiles: Dict[str, ProfileSnapshot] = {}
self._active: Dict[str, float] = {}
self._call_counts: Dict[str, int] = {}
def start(self, segment_id: str):
if segment_id in self._active:
raise RuntimeError(f"Segment {segment_id} already being profiled")
self._active[segment_id] = time.time()
self._call_counts[segment_id] = self._call_counts.get(segment_id, 0) + 1
def stop(self, segment_id: str):
if segment_id not in self._active:
raise RuntimeError(f"Segment {segment_id} not being profiled")
elapsed = time.time() - self._active.pop(segment_id)
count = self._call_counts[segment_id]
if segment_id in self._profiles:
prev = self._profiles[segment_id]
new_elapsed = prev.elapsed + elapsed
new_count = prev.call_count + count
self._profiles[segment_id] = ProfileSnapshot(
elapsed=new_elapsed,
call_count=new_count
)
else:
self._profiles[segment_id] = ProfileSnapshot(
elapsed=elapsed,
call_count=count
)
def get_profile(self, segment_id: str) -> Optional[ProfileSnapshot]:
return self._profiles.get(segment_id)
def get_all_profiles(self) -> Dict[str, ProfileSnapshot]:
return self._profiles.copy()
def reset(self):
self._profiles.clear()
self._active.clear()
self._call_counts.clear()