From 2eca9bb087908fce5fad9e96aec103d9e986a82c Mon Sep 17 00:00:00 2001 From: "Eduardo Sena S. Rosa" Date: Sat, 21 May 2022 00:08:13 -0300 Subject: [PATCH] adding support for multiple threads --- timebudget/timebudget.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/timebudget/timebudget.py b/timebudget/timebudget.py index efb4e1d..043beb8 100644 --- a/timebudget/timebudget.py +++ b/timebudget/timebudget.py @@ -25,6 +25,7 @@ def my_possibly_slow_function(*args): import time from typing import Callable, Optional, Union import warnings +from threading import get_ident __all__ = [ 'timebudget', @@ -67,24 +68,24 @@ def _print(self, msg:str): self.out_stream.flush() def start(self, block_name:str): - if block_name in self.start_times: + if (block_name, get_ident()) in self.start_times: # End should clear out the record, so something odd has happened here. # try/finally should prevent this, but sometimes it doesn't. warnings.warn(f"timebudget is confused: timebudget.start({block_name}) without end") - self.start_times[block_name] = time.time() + self.start_times[(block_name, get_ident())] = time.time() def end(self, block_name:str, quiet:Optional[bool]=None) -> float: """Returns number of ms spent in this block this time. """ if quiet is None: quiet = self.quiet_mode - if block_name not in self.start_times: + if (block_name, get_ident()) not in self.start_times: warnings.warn(f"timebudget is confused: timebudget.end({block_name}) without start") return float('NaN') - elapsed = 1000*(time.time() - self.start_times[block_name]) + elapsed = 1000*(time.time() - self.start_times[(block_name, get_ident())]) self.elapsed_total[block_name] += elapsed self.elapsed_cnt[block_name] += 1 - del self.start_times[block_name] + del self.start_times[(block_name, get_ident())] if not quiet: self._print(f"{block_name} took {ms_format(elapsed)}") return elapsed