-
Notifications
You must be signed in to change notification settings - Fork 0
/
timer.py
67 lines (52 loc) · 2.12 KB
/
timer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import logging
import time
from contextlib import ContextDecorator
from dataclasses import dataclass, field
from typing import Any, Callable, ClassVar
class TimerError(Exception):
"""A custom exception used to report errors in use of Timer class"""
@dataclass
class Timer(ContextDecorator):
"""Time your code using a class, context manager, or decorator"""
timers: ClassVar[dict[str, float]] = dict()
name: None | str = None
text: str = "Elapsed time: {:0.4f} seconds"
logger: None | Callable[[str], None] = logging.info
_start_time: None | float = field(default=None, init=False, repr=False)
def __post_init__(self) -> None:
"""Initialization: add timer to dict of timers"""
if self.name:
self.timers.setdefault(self.name, 0)
def start(self) -> None:
"""Start a new timer"""
if self._start_time is not None:
raise TimerError("Timer is running. Use .stop() to stop it")
self._start_time = time.perf_counter()
def stop(self) -> float:
"""Stop the timer, and report the elapsed time"""
if self._start_time is None:
raise TimerError("Timer is not running. Use .start() to start it")
# Calculate elapsed time
elapsed_time = time.perf_counter() - self._start_time
self._start_time = None
# Report elapsed time only if it is greater than 0.5s
if self.logger and elapsed_time > 0.5:
if self.name:
text = f'{self.name} {self.text.lower()}'
else:
text = self.text
try:
from colorama import Fore
self.logger(f'{Fore.RED}{text.format(elapsed_time)}{Fore.RESET}')
except Exception:
self.logger(f'{text.format(elapsed_time)}')
if self.name:
self.timers[self.name] += elapsed_time
return elapsed_time
def __enter__(self) -> "Timer":
"""Start a new timer as a context manager"""
self.start()
return self
def __exit__(self, *exc_info: Any) -> None:
"""Stop the context manager timer"""
self.stop()