From cddc48c370524b1b7ea9405a9a649f3ad55fec95 Mon Sep 17 00:00:00 2001
From: David Erb <david.erb@diamond.ac.uk>
Date: Thu, 18 May 2023 05:27:18 +0000
Subject: [PATCH] adds profiler

---
 src/dls_utilpack/profiler.py | 58 ++++++++++++++++++++++++++++--------
 1 file changed, 45 insertions(+), 13 deletions(-)

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 = []