| P = ParamSpec('P')
|
| T = TypeVar('T')
|
|
|
|
|
| @dataclass
|
| class TimeContext:
|
| name: str
|
| start_ns: int
|
| end_ns: int | None = None
|
|
|
| _: KW_ONLY
|
|
|
| at_least_ns: int | None = None
|
| _generator: bool = False
|
|
|
| @property
|
| def duration_ns(self) -> int:
|
| if self.end_ns is None:
|
| return timer() - self.start_ns
|
| return self.end_ns - self.start_ns
|
|
|
| @property
|
| def should_report(self) -> bool:
|
| return self.at_least_ns is None or self.duration_ns >= self.at_least_ns
|
|
|
| def close(self) -> None:
|
| if self._generator:
|
| self._generator = False
|
| return
|
| self.end_ns = timer()
|
| if SHOW_TIMINGS and self.should_report:
|
| show_timing(self.name, self.start_ns, self.end_ns)
|
|
|
| def wrap_generator(self, gen: Iterator[T]) -> Iterator[T]:
|
| self._generator = True
|
| try:
|
| yield from gen
|
| finally:
|
| self.close()
|
|
|
|
|
| @contextmanager
|
| def timed_block(name: str, enable: bool = True) -> Iterator[TimeContext]:
|
| ctx = TimeContext(
|
| name, timer(), at_least_ns=100_000_000 if not enable else None
|
| )
|
| try:
|
| yield ctx
|
| finally:
|
| ctx.close()
|
|
|
|
|
| def name_of(fn: Callable[..., Any]) -> str:
|
| if isinstance(fn, functools.partial):
|
| fn = fn.func
|
| return getattr(fn, '__qualname__', repr(fn))
|
|
|
|
|
| def timed(fn: Callable[P, T]) -> Callable[P, T]:
|
| @functools.wraps(fn)
|
| def _wrapper(*args: P.args, **kw: P.kwargs) -> T:
|
| with timed_block(fn.__qualname__) as ctx:
|
| retval = fn(*args, **kw)
|
| if isinstance(retval, GeneratorType):
|
| return ctx.wrap_generator(retval) # type: ignore
|
| else:
|
| return retval
|
| return _wrapper
|
|
|
|
|
| def timed_generator(fn: Callable[P, Iterable[T]]) -> Callable[P, Iterable[T]]:
|
| @functools.wraps(fn)
|
| def _wrapper(*args: P.args, **kw: P.kwargs) -> Iterable[T]:
|
| with timed_block(fn.__qualname__):
|
| yield from fn(*args, **kw)
|
| return _wrapper
|