From 07becf1435e34121596c194b8a2dad61e3b150c7 Mon Sep 17 00:00:00 2001
From: Alex Lubbock <code@alexlubbock.com>
Date: Sat, 17 Aug 2024 01:27:26 +0100
Subject: [PATCH] fix: line_profiler hook

Capture line_profiler data using capturepost_ hook to ensure we
collate it after the wrapped function is executed.
---
 README.md                              |  3 ++-
 microbench/__init__.py                 | 13 ++++++++++---
 microbench/tests/test_line_profiler.py |  1 +
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 7d24a69..cc49eb3 100644
--- a/README.md
+++ b/README.md
@@ -318,7 +318,8 @@ extended example, above.
 
 You can also add functions to your benchmark suite to capture
 extra information at runtime. These functions must be prefixed with `capture_`
-for them to run automatically before the function starts. They take
+for them to run automatically before the function starts, or `capturepost_`
+for them to run automatically when the function completes. They take
 a single argument, `bm_data`, a dictionary to be extended with extra data.
 Care should be taken to avoid overwriting existing key names.
 
diff --git a/microbench/__init__.py b/microbench/__init__.py
index f53cc0a..a131c18 100644
--- a/microbench/__init__.py
+++ b/microbench/__init__.py
@@ -67,11 +67,11 @@ class JSONEncodeWarning(Warning):
 
 _UNENCODABLE_PLACEHOLDER_VALUE = '__unencodable_as_json__'
 
+
 class MicroBench(object):
     def __init__(self, outfile=None, json_encoder=JSONEncoder,
                  tz=timezone.utc, iterations=1,
                  *args, **kwargs):
-        self._capture_before = []
         if args:
             raise ValueError('Only keyword arguments are allowed')
         self._bm_static = kwargs
@@ -109,7 +109,7 @@ def pre_start_triggers(self, bm_data):
         for method_name in dir(self):
             if method_name.startswith('capture_'):
                 method = getattr(self, method_name)
-                if callable(method) and method not in self._capture_before:
+                if callable(method):
                     method(bm_data)
 
         # Initialise telemetry thread
@@ -132,6 +132,13 @@ def post_finish_triggers(self, bm_data):
             timeout = getattr(self, 'telemetry_timeout', 30)
             self._telemetry_thread.join(timeout)
 
+        # Run capturepost triggers
+        for method_name in dir(self):
+            if method_name.startswith('capturepost_'):
+                method = getattr(self, method_name)
+                if callable(method):
+                    method(bm_data)
+
     def pre_run_triggers(self, bm_data):
         bm_data['_run_start'] = datetime.now(self.tz)
 
@@ -345,7 +352,7 @@ class MBLineProfiler(object):
     slightly slow down the execution of your function, so it's not recommended
     in production.
     """
-    def capture_line_profile(self, bm_data):
+    def capturepost_line_profile(self, bm_data):
         bm_data['line_profiler'] = base64.encodebytes(
             pickle.dumps(self._line_profiler.get_stats())
         ).decode('utf8')
diff --git a/microbench/tests/test_line_profiler.py b/microbench/tests/test_line_profiler.py
index 8c9c39d..f21410f 100644
--- a/microbench/tests/test_line_profiler.py
+++ b/microbench/tests/test_line_profiler.py
@@ -24,3 +24,4 @@ def my_function():
     lp = MBLineProfiler.decode_line_profile(results['line_profiler'][0])
     assert lp.__class__.__name__ == 'LineStats'
     MBLineProfiler.print_line_profile(results['line_profiler'][0])
+    assert not all(len(v) == 0 for v in lp.timings.values()), "No timings present"