diff --git a/src/dls_utilpack/profiler.py b/src/dls_utilpack/profiler.py index 983aa1a..0c5b9b8 100644 --- a/src/dls_utilpack/profiler.py +++ b/src/dls_utilpack/profiler.py @@ -6,7 +6,7 @@ logger = logging.getLogger(__name__) -class Profile: +class Context: """ Class that holds the total execution time and call count of potentially multiple executions of the same label. Reports its contents as a single-line string. @@ -14,18 +14,34 @@ class Profile: TODO: Make profiles nestable by accumulating them at the context close. """ - def __init__(self, label): + def __init__(self, label: str, profiler: Profiler): + self.__profiler = profiler + self.__label = label self.__start_time = 0 - self.__seconds = 0.0 - self.__count = 0 def __enter__(self): self.__start_time = time.time() def __exit__(self, exc_type, exc_val, exc_tb): - self.__seconds = time.time() - self.__start_time - self.__count += 1 + self.__profiler.accumulate( + self.__label, + time.time() - self.__start_time, + ) + + +class Profile: + """ + Class that holds the total execution time and call count of potentially multiple executions of the same label. + Reports its contents as a single-line string. + + TODO: Make profiles nestable by accumulating them at the context close. + """ + + def __init__(self, label): + self.__label = label + self.__seconds = 0.0 + self.__count = 0 def __str__(self): if self.__count == 0: @@ -47,10 +63,11 @@ class Profiler: def __init__(self): self.__profiles = {} + self.__lock = threading.RLock() - def profile(self, label: str) -> Profile: + def profile(self, label: str) -> Context: """ - Return the profile for the given label. Uses previously existing profile, if any, or makes a new instance. + Return a context to hold the profile timing. Args: label (str): label identifying the profile @@ -58,12 +75,27 @@ def profile(self, label: str) -> Profile: Returns: Profile: a new profile object, or previously existing one """ - profile = self.__profiles.get(label) - if profile is None: - profile = Profile(label) - self.__profiles[label] = profile + return Context(label, self) + + def accumulate(self, label: str, seconds: float) -> None: + """ + Accumulate a report into the profile for the given label. + Uses previously existing profile, if any, or makes a new instance. + + Args: + label (str): label identifying the profile - return profile + Returns: + Profile: a new profile object, or previously existing one + """ + with self.__lock: + profile = self.__profiles.get(label) + if profile is None: + profile = Profile(label) + self.__profiles[label] = profile + + profile.seconds += seconds + profile.count += 1 def __str__(self) -> str: lines = []