fix: line_profiler hook

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.
 

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')

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"