Add support for teardown handler to pedantic mode

Sometimes benchmarks have side effects which need to be cleaned up after every
round. For example if a benchmark writes to a file, then you might want to
delete it in between rounds and start from a clean slate.

It's already possible to pass a setup function to pedantic mode, this PR
introduces a similar mechanism but for cleaning up resources after a round has
been competed by passing a cleanup function.

src/pytest_benchmark/fixture.py

@@ -143,14 +143,14 @@ def __call__(self, function_to_benchmark, *args, **kwargs):
             self.has_error = True
             raise
 
-    def pedantic(self, target, args=(), kwargs=None, setup=None, rounds=1, warmup_rounds=0, iterations=1):
+    def pedantic(self, target, args=(), kwargs=None, setup=None, cleanup=None, rounds=1, warmup_rounds=0, iterations=1):
         if self._mode:
             self.has_error = True
             raise FixtureAlreadyUsed('Fixture can only be used once. Previously it was used in %s mode.' % self._mode)
         try:
             self._mode = 'benchmark.pedantic(...)'
             return self._raw_pedantic(
-                target, args=args, kwargs=kwargs, setup=setup, rounds=rounds, warmup_rounds=warmup_rounds, iterations=iterations
+                target, args=args, kwargs=kwargs, setup=setup, cleanup=cleanup, rounds=rounds, warmup_rounds=warmup_rounds, iterations=iterations
             )
         except Exception:
             self.has_error = True
@@ -187,7 +187,7 @@ def _raw(self, function_to_benchmark, *args, **kwargs):
             function_result = function_to_benchmark(*args, **kwargs)
         return function_result
 
-    def _raw_pedantic(self, target, args=(), kwargs=None, setup=None, rounds=1, warmup_rounds=0, iterations=1):
+    def _raw_pedantic(self, target, args=(), kwargs=None, setup=None, cleanup=None, rounds=1, warmup_rounds=0, iterations=1):
         if kwargs is None:
             kwargs = {}
 
@@ -226,6 +226,9 @@ def make_arguments(args=args, kwargs=kwargs):
             runner = self._make_runner(target, args, kwargs)
             runner(loops_range)
 
+            if cleanup is not None:
+                cleanup()
+
         for _ in range(rounds):
             args, kwargs = make_arguments()
 
@@ -236,6 +239,11 @@ def make_arguments(args=args, kwargs=kwargs):
                 duration, result = runner(loops_range)
             stats.update(duration)
 
+            if cleanup is not None:
+                cleanup()
+
+
+        # TODO: Figure out what this is about
         if loops_range:
             args, kwargs = make_arguments()
             result = target(*args, **kwargs)
@@ -246,6 +254,9 @@ def make_arguments(args=args, kwargs=kwargs):
             profile.runcall(target, *args, **kwargs)
             self.stats.cprofile_stats = pstats.Stats(profile)
 
+            if cleanup is not None:
+                cleanup()
+
         return result
 
     def weave(self, target, **kwargs):

tests/test_pedantic.py

@@ -22,6 +22,19 @@ def setup():
     assert runs == [(1, 2)]
 
 
+def test_cleanup(benchmark):
+    runs = []
+
+    def stuff():
+        runs.append("stuff")
+
+    def cleanup():
+        runs.append("cleanup")
+
+    benchmark.pedantic(stuff, cleanup=cleanup)
+    assert runs == ["stuff", "cleanup"]
+
+
 @pytest.mark.benchmark(cprofile=True)
 def test_setup_cprofile(benchmark):
     runs = []
@@ -36,6 +49,22 @@ def setup():
     assert runs == [(1, 2), (1, 2)]
 
 
+@pytest.mark.benchmark(cprofile=True)
+def test_cleanup_cprofile(benchmark):
+    runs = []
+
+    def stuff():
+        runs.append("stuff")
+
+    def cleanup():
+        runs.append("cleanup")
+
+    benchmark.pedantic(stuff, cleanup=cleanup)
+    assert runs == ["stuff", "cleanup", "stuff", "cleanup"]
+
+    runs = []
+
+
 def test_args_kwargs(benchmark):
     runs = []
 
@@ -100,6 +129,18 @@ def setup():
     benchmark.pedantic(stuff, setup=setup, rounds=10)
     assert runs == [(1, 2)] * 10
 
+def test_cleanup_many_rounds(benchmark):
+    runs = []
+
+    def stuff():
+        runs.append("stuff")
+
+    def cleanup():
+        runs.append("cleanup")
+
+    benchmark.pedantic(stuff, cleanup=cleanup, rounds=10)
+    assert runs == ["stuff", "cleanup"] * 10
+
 
 def test_cant_use_both_args_and_setup_with_return(benchmark):
     runs = []