Refactor: Improve time extraction and testability of Metrics class

tools/metrics.py

@@ -1,13 +1,19 @@
+# tools/metrics.py
 import cProfile
 import pstats
-import io
+import io 
 from functools import wraps
 from collections import defaultdict
+import logging
 
 class Metrics:
-    def __init__(self):
+    def __init__(self, logger=None):
         self.call_count = defaultdict(int)
         self.total_time = defaultdict(float)
+        self.logger = logger if logger else logging.getLogger(__name__)
+        if not self.logger.handlers:
+            self.logger.addHandler(logging.NullHandler())
+        self.logger.debug("Metrics instance initialized.")
 
     def measure(self, func):
         @wraps(func)
@@ -16,30 +22,58 @@ class Metrics:
             
             pr = cProfile.Profile()
             pr.enable()
-            
             result = func(*args, **kwargs)
-            
             pr.disable()
-            s = io.StringIO()
-            ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
-            ps.print_stats()
             
-            # Extract the total time spent in the function
-            time_spent = float(s.getvalue().split('\n')[0].split()[-2])
-            self.total_time[func.__name__] += time_spent
+            ps = pstats.Stats(pr)
+            
+            func_code = func.__code__
+            func_key_tuple = (func_code.co_filename, func_code.co_firstlineno, func_code.co_name)
+            
+            time_spent_for_func = 0.0
+            if func_key_tuple in ps.stats:
+                time_spent_for_func = ps.stats[func_key_tuple][3] # [3] is cumulative time (ct)
+            else:
+                # Fallback: try to find by function name if exact key fails (e.g. due to decorators changing code object details slightly)
+                # This is less precise and might pick up other functions if names are not unique across files.
+                found_by_name = False
+                for key, stat in ps.stats.items():
+                    if key[2] == func.__name__: # key[2] is function name
+                        time_spent_for_func = stat[3] # cumulative time
+                        self.logger.debug(f"Found stats for {func.__name__} by name {key} after primary key failed.")
+                        found_by_name = True
+                        break
+                if not found_by_name:
+                    self.logger.warning(
+                        f"Could not find exact cProfile stats for {func.__name__} with key {func_key_tuple} or by name. "
+                        f"Time for this call will be recorded as 0. This might occur for non-Python functions or due to complex decorators."
+                    )
+            
+            self.total_time[func.__name__] += time_spent_for_func
+            self.logger.debug(f"Measured cumulative time for {func.__name__}: {time_spent_for_func:.6f}s")
             
             return result
         return wrapper
 
     def get_metrics(self):
-        metrics = {}
+        metrics_data = {}
         for func_name in self.call_count:
-            metrics[func_name] = {
-                'call_count': self.call_count[func_name],
-                'total_time': self.total_time[func_name],
-                'average_time': self.total_time[func_name] / self.call_count[func_name] if self.call_count[func_name] > 0 else 0
+            count = self.call_count[func_name]
+            total_t = self.total_time[func_name]
+            metrics_data[func_name] = {
+                'call_count': count,
+                'total_time': round(total_t, 6),
+                'average_time': round(total_t / count, 6) if count > 0 else 0
             }
-        return metrics
+        return metrics_data
 
-# Create a global instance of Metrics
-metrics = Metrics()
+    def clear_metrics(self):
+        self.call_count.clear()
+        self.total_time.clear()
+        self.logger.info("Metrics cleared.")
+
+# Global instance for convenience
+_metrics_instance_logger = logging.getLogger(__name__ + ".global_instance")
+if not _metrics_instance_logger.handlers:
+    _metrics_instance_logger.addHandler(logging.NullHandler())
+metrics = Metrics(logger=_metrics_instance_logger)