Add ability to register hooks from code (#236)

Add hooks parameter to Runner.

Co-authored-by: Victor Stinner <vstinner@python.org>

Author: sergey-miryanov

doc/api.rst

@@ -485,10 +485,39 @@ BenchmarkSuite class
       It can be ``None``.
 
 
+HookBase class
+--------------
+
+.. class:: HookBase()
+
+   Hook used to do some actions before and after benchmark's run. It can collect
+   statistics, run external tools.
+
+   Methods:
+
+   .. method:: __enter__()
+
+      Called immediately before runnig benchmark code.
+
+      May be called multiple times per instance.
+
+   .. method:: __exit__(exc_type, exc_value, traceback):
+
+      Called immediately after running benchmark code.
+
+      May be called multiple times per instance.
+
+   .. method:: teardown(metadata)
+
+      Called when the hook is completed for a process.
+
+      May add any information collected to the passed-in ``metadata`` dictionary.
+
+
 Runner class
 ------------
 
-.. class:: Runner(values=3, warmups=1, processes=20, loops=0, min_time=0.1, metadata=None, show_name=True, program_args=None, add_cmdline_args=None)
+.. class:: Runner(values=3, warmups=1, processes=20, loops=0, min_time=0.1, metadata=None, show_name=True, program_args=None, add_cmdline_args=None, hooks=None)
 
    Tool to run a benchmark in text mode.
 
@@ -512,6 +541,10 @@ Runner class
    (``list``) which must be modified in place and *args* is the :attr:`args`
    attribute of the runner.
 
+   *hooks* is a list or tuple of custom hooks. Each hook should implement
+   :class:`HookBase` and have ``name`` attribute and ``load()`` method that
+   returns instance of the hook.
+
    If *show_name* is true, displays the benchmark name.
 
    If isolated CPUs are detected, the CPU affinity is automatically

pyperf/_hooks.py

@@ -29,18 +29,18 @@ def get_hook_names():
     return (x.name for x in get_hooks())
 
 
-def get_selected_hooks(hook_names):
+def get_selected_hooks(hook_names, hooks=None):
     if hook_names is None:
         return
 
-    hook_mapping = {hook.name: hook for hook in get_hooks()}
+    hook_mapping = {hook.name: hook for hook in get_hooks() + tuple(hooks or ())}
     for hook_name in hook_names:
         yield hook_mapping[hook_name]
 
 
-def instantiate_selected_hooks(hook_names):
+def instantiate_selected_hooks(hook_names, hooks=None):
     hook_managers = {}
-    for hook in get_selected_hooks(hook_names):
+    for hook in get_selected_hooks(hook_names, hooks):
         try:
             hook_managers[hook.name] = hook.load()()
         except HookError as e:

pyperf/_runner.py

@@ -77,7 +77,7 @@ def __init__(self, values=None, processes=None,
                  loops=0, min_time=0.1, metadata=None,
                  show_name=True,
                  program_args=None, add_cmdline_args=None,
-                 _argparser=None, warmups=1):
+                 _argparser=None, warmups=1, hooks=None):
 
         # Watchdog: ensure that only once instance of Runner (or a Runner
         # subclass) is created per process to prevent bad surprises
@@ -247,7 +247,11 @@ def __init__(self, values=None, processes=None,
                             help='Collect profile data using cProfile '
                                  'and output to the given file.')
 
-        hook_names = list(get_hook_names())
+        self._custom_hooks = hooks or ()
+        if not isinstance(self._custom_hooks, (list, tuple)):
+            raise RuntimeError("hooks parameter should be list, tuple or None")
+
+        hook_names = list(get_hook_names()) + [hook.name for hook in self._custom_hooks]
         parser.add_argument(
             '--hook', action="append", choices=hook_names,
             metavar=f"{', '.join(x for x in hook_names if not x.startswith('_'))}",

pyperf/_worker.py

@@ -31,6 +31,7 @@ def __init__(self, runner, name, task_func, func_metadata):
         self.args = args
         self.task_func = task_func
         self.loops = args.loops
+        self._custom_hooks = runner._custom_hooks
 
         self.metadata = dict(runner.metadata)
         if func_metadata:
@@ -61,7 +62,7 @@ def _compute_values(self, values, nvalue,
 
         task_func = self.task_func
 
-        hook_managers = instantiate_selected_hooks(args.hook)
+        hook_managers = instantiate_selected_hooks(args.hook, self._custom_hooks)
         if len(hook_managers):
             self.metadata["hooks"] = ", ".join(hook_managers.keys())
 

pyperf/tests/test_runner.py

@@ -10,6 +10,7 @@
 
 import pyperf
 from pyperf import tests
+from pyperf._hooks import HookBase
 from pyperf._utils import create_pipe, MS_WINDOWS, shell_quote
 
 
@@ -518,6 +519,46 @@ def test_hook_command(self):
         self.assertEqual(bench.get_metadata()["hooks"],
                          "_test_hook")
 
+    def test_custom_hook(self):
+        class State:
+            inited = 0
+            entered = 0
+            exited = 0
+
+        s = State()
+
+        class CustomHook(HookBase):
+            name = "custom_hook"
+
+            @staticmethod
+            def load():
+                return lambda: CustomHook(s)
+
+            def __init__(self, state):
+                self.state = state
+                self.state.inited += 1
+
+            def __enter__(self):
+                self.state.entered += 1
+
+            def __exit__(self, _exc_type, _exc_value, _traceback):
+                self.state.exited += 1
+
+            def teardown(self, metadata):
+                metadata[self.name] = "done"
+
+        args = '-l1 -w0 -n1 --worker --verbose --hook custom_hook'.split()
+        runner = self.create_runner(args, hooks=[CustomHook])
+        def time_func(loops):
+            return 1.0
+
+        bench = runner.bench_time_func('bench1', time_func)
+        self.assertEqual(bench.get_metadata()["hooks"], CustomHook.name)
+        self.assertEqual(bench.get_metadata()[CustomHook.name], "done")
+        self.assertEqual(s.inited, 1)
+        self.assertEqual(s.entered, 1)
+        self.assertEqual(s.exited, 1)
+
     def test_single_instance(self):
         runner1 = self.create_runner([])   # noqa
         with self.assertRaises(RuntimeError):