56 lines
1.7 KiB
Python
Executable File
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()
|