67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
from typing import Union, Optional, Any, TypeVar
|
|
from collections.abc import Iterable, Collection, Sequence, Callable, Generator, AsyncGenerator
|
|
|
|
import time
|
|
|
|
class Timer():
|
|
def __init__(self,
|
|
print_format: Optional[Union[str, bool]] = None,
|
|
start: bool = False,
|
|
run_fn: Optional[Callable[[float], Any]] = None,
|
|
run_fn_full: Optional[Callable[[float, float, float], Any]] = None,
|
|
print_fn: Callable[[str], Any] = print,
|
|
time_fn: Callable[[], float] = time.monotonic,
|
|
):
|
|
if isinstance(print_format, bool):
|
|
if print_format is True:
|
|
print_format = "{secs}s"
|
|
else:
|
|
print_format = None
|
|
|
|
self.print_format = print_format
|
|
self.run_fn = run_fn
|
|
self.run_fn_full = run_fn_full
|
|
self.print_fn = print_fn
|
|
self.time_fn = time_fn
|
|
|
|
# start/end timestamps, by default monotonic time unless time_fn is set
|
|
self.start_ts = None
|
|
self.end_ts = None
|
|
self.seconds: Optional[float] = None
|
|
|
|
if start:
|
|
self.start()
|
|
|
|
def start(self):
|
|
self.start_ts = self.time_fn()
|
|
|
|
def end(self, run: bool = True) -> float:
|
|
self.end_ts = self.time_fn()
|
|
self.seconds = self.end_ts - self.start_ts
|
|
|
|
if run:
|
|
self._end()
|
|
|
|
return self.seconds
|
|
|
|
def _end(self):
|
|
if self.print_format is not None:
|
|
self.print_fn(self.print_format.format(
|
|
self.seconds, self.seconds, self.seconds, self.seconds, self.seconds, self.seconds,
|
|
s = self.seconds, secs = self.seconds, seconds = self.seconds,
|
|
start = self.start_ts, end = self.end_ts,
|
|
))
|
|
|
|
if self.run_fn is not None:
|
|
self.run_fn(self.seconds)
|
|
|
|
if self.run_fn_full is not None:
|
|
self.run_fn_full(self.start_ts, self.end_ts, self.seconds)
|
|
|
|
def __enter__(self):
|
|
self.start()
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
self.end()
|