| 150 | |
| 151 | |
| 152 | class timeit: |
| 153 | _history: Dict[str, List['timeit']] = {} |
| 154 | |
| 155 | def __init__(self, name: str = None, verbose: bool = True, average: bool = False): |
| 156 | self.name = name |
| 157 | self.verbose = verbose |
| 158 | self.start = None |
| 159 | self.end = None |
| 160 | self.average = average |
| 161 | if average and name not in timeit._history: |
| 162 | timeit._history[name] = [] |
| 163 | |
| 164 | def __call__(self, func: Callable): |
| 165 | import inspect |
| 166 | if inspect.iscoroutinefunction(func): |
| 167 | async def wrapper(*args, **kwargs): |
| 168 | with timeit(self.name or func.__qualname__): |
| 169 | ret = await func(*args, **kwargs) |
| 170 | return ret |
| 171 | return wrapper |
| 172 | else: |
| 173 | def wrapper(*args, **kwargs): |
| 174 | with timeit(self.name or func.__qualname__): |
| 175 | ret = func(*args, **kwargs) |
| 176 | return ret |
| 177 | return wrapper |
| 178 | |
| 179 | def __enter__(self): |
| 180 | self.start = time.time() |
| 181 | return self |
| 182 | |
| 183 | @property |
| 184 | def time(self) -> float: |
| 185 | assert self.start is not None, "Time not yet started." |
| 186 | assert self.end is not None, "Time not yet ended." |
| 187 | return self.end - self.start |
| 188 | |
| 189 | @property |
| 190 | def average_time(self) -> float: |
| 191 | assert self.average, "Average time not available." |
| 192 | return sum(t.time for t in timeit._history[self.name]) / len(timeit._history[self.name]) |
| 193 | |
| 194 | @property |
| 195 | def history(self) -> List['timeit']: |
| 196 | return timeit._history.get(self.name, []) |
| 197 | |
| 198 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 199 | self.end = time.time() |
| 200 | if self.average: |
| 201 | timeit._history[self.name].append(self) |
| 202 | if self.verbose: |
| 203 | if self.average: |
| 204 | avg = self.average_time |
| 205 | print(f"{self.name or 'It'} took {avg:.6f} seconds in average.") |
| 206 | else: |
| 207 | print(f"{self.name or 'It'} took {self.time:.6f} seconds.") |
| 208 | |
| 209 | |