64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
|
import time
|
||
|
|
from typing import Optional, List, Dict, Any
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ExecutionTrace:
|
||
|
|
segment_id: str
|
||
|
|
byte_start: int
|
||
|
|
byte_end: int
|
||
|
|
duration: float
|
||
|
|
exception: Optional[str] = None
|
||
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
|
||
|
|
class ExecutionTracer:
|
||
|
|
def __init__(self, verbose: bool = False):
|
||
|
|
self.verbose = verbose
|
||
|
|
self.traces: List[ExecutionTrace] = []
|
||
|
|
|
||
|
|
def trace_segment(self, segment_id: str, byte_start: int, byte_end: int):
|
||
|
|
return _SegmentTraceContext(self, segment_id, byte_start, byte_end)
|
||
|
|
|
||
|
|
def record_trace(self, trace: ExecutionTrace):
|
||
|
|
self.traces.append(trace)
|
||
|
|
if self.verbose:
|
||
|
|
print(f"[TRACE] {trace.segment_id} [{trace.byte_start}:{trace.byte_end}] {trace.duration:.4f}s")
|
||
|
|
|
||
|
|
def get_traces(self) -> List[ExecutionTrace]:
|
||
|
|
return self.traces.copy()
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
self.traces.clear()
|
||
|
|
|
||
|
|
|
||
|
|
class _SegmentTraceContext:
|
||
|
|
def __init__(self, tracer: ExecutionTracer, segment_id: str, byte_start: int, byte_end: int):
|
||
|
|
self.tracer = tracer
|
||
|
|
self.segment_id = segment_id
|
||
|
|
self.byte_start = byte_start
|
||
|
|
self.byte_end = byte_end
|
||
|
|
self.start_time = None
|
||
|
|
self.exception = None
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
self.start_time = time.time()
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
|
|
duration = time.time() - self.start_time
|
||
|
|
exception = None
|
||
|
|
if exc_type:
|
||
|
|
exception = f"{exc_type.__name__}: {exc_val}"
|
||
|
|
|
||
|
|
trace = ExecutionTrace(
|
||
|
|
segment_id=self.segment_id,
|
||
|
|
byte_start=self.byte_start,
|
||
|
|
byte_end=self.byte_end,
|
||
|
|
duration=duration,
|
||
|
|
exception=exception
|
||
|
|
)
|
||
|
|
self.tracer.record_trace(trace)
|
||
|
|
return False
|