Repository: dwavesystems/dwave-cloud-client

Issue:
Event dispatch decorator
Call `dispatch_event` before and after decorated function execution, automating args and rval collection/reporting.

Comments:

Patch:
diff --git a/dwave/cloud/client.py b/dwave/cloud/client.py
index 7986b84..db3c444 100644
--- a/dwave/cloud/client.py
+++ b/dwave/cloud/client.py
@@ -73,7 +73,7 @@ from dwave.cloud.config import load_config, parse_float, parse_int, parse_boolea
 from dwave.cloud.solver import Solver, available_solvers
 from dwave.cloud.concurrency import PriorityThreadPoolExecutor
 from dwave.cloud.upload import ChunkedData
-from dwave.cloud.events import dispatch_event
+from dwave.cloud.events import dispatches_events
 from dwave.cloud.utils import (
     TimeoutingHTTPAdapter, BaseUrlSession, user_agent,
     datetime_to_timestamp, utcnow, epochnow, cached, retried, is_caused_by)
@@ -343,6 +343,7 @@ class Client(object):
         logger.debug("Creating %s.Client() with: %r", _client, config)
         return _clients[_client](**config)
 
+    @dispatches_events('client_init')
     def __init__(self, endpoint=None, token=None, solver=None, **kwargs):
         # for (reasonable) backwards compatibility, accept only the first few
         # positional args.
@@ -354,9 +355,6 @@ class Client(object):
         if solver is not None:
             kwargs.setdefault('solver', solver)
 
-        dispatched_args = kwargs.copy()
-        dispatch_event('before_client_init', obj=self, args=dispatched_args)
-
         logger.debug("Client init called with: %r", kwargs)
 
         # derive instance-level defaults from class defaults and init defaults
@@ -523,9 +521,6 @@ class Client(object):
         self._encode_problem_executor = \
             ThreadPoolExecutor(self._ENCODE_PROBLEM_THREAD_COUNT)
 
-        dispatch_event(
-            'after_client_init', obj=self, args=dispatched_args, return_value=None)
-
     def create_session(self):
         """Create a new requests session based on client's (self) params.
 
@@ -728,6 +723,7 @@ class Client(object):
         self._load(future)
         return future
 
+    @dispatches_events('get_solvers')
     def get_solvers(self, refresh=False, order_by='avg_load', **filters):
         """Return a filtered list of solvers handled by this client.
 
@@ -913,9 +909,6 @@ class Client(object):
             )
         """
 
-        args = dict(refresh=refresh, order_by=order_by, filters=filters)
-        dispatch_event('before_get_solvers', obj=self, args=args)
-
         def covers_op(prop, val):
             """Does LHS `prop` (range) fully cover RHS `val` (range or item)?"""
 
@@ -1064,9 +1057,6 @@ class Client(object):
         if sort_reverse:
             solvers.reverse()
 
-        dispatch_event(
-            'after_get_solvers', obj=self, args=args, return_value=solvers)
-
         return solvers
 
     def solvers(self, refresh=False, **filters):
diff --git a/dwave/cloud/events.py b/dwave/cloud/events.py
index 8e9b53d..66a0a39 100644
--- a/dwave/cloud/events.py
+++ b/dwave/cloud/events.py
@@ -19,6 +19,8 @@ inspection.
 """
 
 import logging
+import inspect
+from functools import wraps
 
 __all__ = ['add_handler']
 
@@ -41,13 +43,17 @@ _client_event_hooks_registry = {
 def add_handler(name, handler):
     """Register a `handler` function to be called on event `name`.
 
-    Handler's signature are::
+    Handler signatures are::
 
         def before_event_handler(event_name, obj, args):
-            pass
+            # called just before `obj.method(**args)` executes
 
         def after_event_handler(event_name, obj, args, return_value):
-            pass
+            # function succeeded with `return_value`
+
+        def after_event_handler(event_name, obj, args, exception):
+            # function failed with `exception` raised
+            # after event handler invocation, exception is re-raised
 
     """
 
@@ -73,3 +79,35 @@ def dispatch_event(name, *args, **kwargs):
         except Exception as e:
             logger.debug("Exception in {!r} event handler {!r}: {!r}".format(
                 name, handler, e))
+
+
+class dispatches_events:
+    """Decorate function to :func:`.dispatch_event` on entry and exit."""
+
+    def __init__(self, basename):
+        self.before_eventname = 'before_' + basename
+        self.after_eventname = 'after_' + basename
+
+    def __call__(self, fn):
+        if not callable(fn):
+            raise TypeError("decorated object must be callable")
+
+        @wraps(fn)
+        def wrapped(*pargs, **kwargs):
+            sig = inspect.signature(fn)
+            bound = sig.bind(*pargs, **kwargs)
+            bound.apply_defaults()
+            args = bound.arguments
+            obj = args.pop('self', None)
+
+            dispatch_event(self.before_eventname, obj=obj, args=args)
+            try:
+                rval = fn(*pargs, **kwargs)
+            except Exception as exc:
+                dispatch_event(self.after_eventname, obj=obj, args=args, exception=exc)
+                raise
+            else:
+                dispatch_event(self.after_eventname, obj=obj, args=args, return_value=rval)
+                return rval
+
+        return wrapped
diff --git a/dwave/cloud/solver.py b/dwave/cloud/solver.py
index 56b7f34..fbaf0a8 100644
--- a/dwave/cloud/solver.py
+++ b/dwave/cloud/solver.py
@@ -40,7 +40,7 @@ from dwave.cloud.coders import (
 from dwave.cloud.utils import uniform_iterator, reformat_qubo_as_ising
 from dwave.cloud.computation import Future
 from dwave.cloud.concurrency import Present
-from dwave.cloud.events import dispatch_event
+from dwave.cloud.events import dispatches_events
 
 # Use numpy if available for fast encoding/decoding
 try:
@@ -404,6 +404,7 @@ class BaseUnstructuredSolver(BaseSolver):
 
         return body_data
 
+    @dispatches_events('sample')
     def sample_problem(self, problem, problem_type=None, label=None, **params):
         """Sample from the specified problem.
 
@@ -877,6 +878,7 @@ class StructuredSolver(BaseSolver):
         return self._sample(problem_type, bqm.linear, bqm.quadratic, bqm.offset,
                             params, label=label, undirected_biases=True)
 
+    @dispatches_events('sample')
     def _sample(self, type_, linear, quadratic, offset, params,
                 label=None, undirected_biases=False):
         """Internal method for `sample_ising`, `sample_qubo` and `sample_bqm`.
@@ -906,11 +908,6 @@ class StructuredSolver(BaseSolver):
             :class:`~dwave.cloud.computation.Future`
         """
 
-        args = dict(type_=type_, linear=linear, quadratic=quadratic,
-                    offset=offset, params=params, label=label,
-                    undirected_biases=undirected_biases)
-        dispatch_event('before_sample', obj=self, args=args)
-
         # Check the problem
         if not self.check_problem(linear, quadratic):
             raise InvalidProblemError("Problem graph incompatible with solver.")
@@ -948,8 +945,6 @@ class StructuredSolver(BaseSolver):
         logger.debug("Submitting new problem to: %s", self.id)
         self.client._submit(body, computation)
 
-        dispatch_event('after_sample', obj=self, args=args, return_value=computation)
-
         return computation
 
     def _format_params(self, type_, params):


Test patch:
diff --git a/tests/test_events.py b/tests/test_events.py
index 76a933c..ca95168 100644
--- a/tests/test_events.py
+++ b/tests/test_events.py
@@ -14,21 +14,28 @@
 
 import unittest
 
+try:
+    import dimod
+except ImportError:
+    dimod = None
+
 from dwave.cloud.client import Client
-from dwave.cloud.solver import Solver
-from dwave.cloud.events import add_handler
+from dwave.cloud.solver import StructuredSolver, UnstructuredSolver
+from dwave.cloud.events import add_handler, dispatches_events
+from dwave.cloud.concurrency import Present
 
 
 class TestEventDispatch(unittest.TestCase):
 
     def setUp(self):
         # mock client
-        self.client = Client(token='token', solver={'qpu': True})
+        self.client = Client(token='token', solver=dict(name__contains='test'))
         self.client._fetch_solvers = lambda **kw: self.solvers
         self.client._submit = lambda *pa, **kw: None
+        self.client.upload_problem_encoded = lambda *pa, **kw: Present(result=mock_problem_id)
 
         # mock solvers
-        self.solver = Solver(client=self.client, data={
+        self.structured_solver = StructuredSolver(client=self.client, data={
             "properties": {
                 "supported_problem_types": ["qubo", "ising"],
                 "qubits": [0, 1, 2],
@@ -46,11 +53,24 @@ class TestEventDispatch(unittest.TestCase):
                 "category": "qpu",
                 "tags": ["lower_noise"]
             },
-            "id": "solver1",
+            "id": "test-qpu-solver",
             "description": "A test solver 1",
             "status": "online"
         })
-        self.solvers = [self.solver]
+        self.unstructured_solver = UnstructuredSolver(client=self.client, data={
+            "properties": {
+                "supported_problem_types": ["bqm"],
+                "parameters": {"num_reads": "Number of samples to return."},
+                "category": "hybrid",
+            },
+            "id": "test-unstructured-solver",
+            "description": "A test unstructured solver"
+        })
+        self.solvers = [self.structured_solver, self.unstructured_solver]
+
+        # reset all event handlers
+        from dwave.cloud.events import _client_event_hooks_registry as reg
+        reg.update({k: [] for k in reg})
 
     def test_validation(self):
         """Event name and handler are validated."""
@@ -77,16 +97,18 @@ class TestEventDispatch(unittest.TestCase):
         # test entry values
         before = memo['before_client_init']
         self.assertEqual(before['obj'], client)
-        self.assertNotIn('endpoint', before['args'])
+        self.assertIn('endpoint', before['args'])
         self.assertIn('token', before['args'])
+        self.assertIn('kwargs', before['args'])
         self.assertEqual(before['args']['token'], 'token')
-        self.assertEqual(before['args']['unknown'], 'unknown')
+        self.assertEqual(before['args']['kwargs']['unknown'], 'unknown')
 
         # test exit values
         after = memo['after_client_init']
         self.assertEqual(after['obj'], client)
         self.assertEqual(after['args']['token'], 'token')
-        self.assertEqual(after['args']['unknown'], 'unknown')
+        self.assertEqual(after['args']['kwargs']['unknown'], 'unknown')
+        self.assertEqual(after['args']['endpoint'], None)
         self.assertEqual(after['return_value'], None)
 
     def test_get_solvers(self):
@@ -108,17 +130,15 @@ class TestEventDispatch(unittest.TestCase):
         self.assertEqual(before['obj'], self.client)
         self.assertIn('refresh', before['args'])
         self.assertIn('filters', before['args'])
-        self.assertIn('qpu', before['args']['filters'])
+        self.assertIn('name__contains', before['args']['filters'])
 
         # test exit values
         after = memo['after_get_solvers']
         self.assertEqual(after['obj'], self.client)
-        self.assertIn('qpu', after['args']['filters'])
+        self.assertIn('name__contains', after['args']['filters'])
         self.assertEqual(after['return_value'], self.solvers)
 
-    def test_sample(self):
-        """Before/After solver sample events are dispatched with correct signatures."""
-
+    def subtest_sample(self, solver):
         # setup event handlers
         memo = {}
         def handler(event, **data):
@@ -132,18 +152,75 @@ class TestEventDispatch(unittest.TestCase):
         quad = {(0, 1): 1}
         offset = 2
         params = dict(num_reads=100)
-        future = self.solver.sample_ising(lin, quad, offset, **params)
+        future = solver.sample_ising(lin, quad, offset, **params)
 
         # test entry values
         before = memo['before_sample']
-        args = dict(type_='ising', linear=lin, quadratic=quad,
-                    offset=offset, params=params,
-                    undirected_biases=False, label=None)
-        self.assertEqual(before['obj'], self.solver)
+        if solver.qpu:
+            args = dict(type_='ising', linear=lin, quadratic=quad,
+                        offset=offset, params=params,
+                        undirected_biases=False, label=None)
+        elif solver.hybrid:
+            if not dimod:
+                self.skipTest("dimod not installed")
+            bqm = dimod.BQM.from_ising(lin, quad, offset)
+            args = dict(problem=bqm, problem_type=None, label=None, params=params)
+
+        self.assertEqual(before['obj'], solver)
         self.assertDictEqual(before['args'], args)
 
         # test exit values
         after = memo['after_sample']
-        self.assertEqual(after['obj'], self.solver)
+        self.assertEqual(after['obj'], solver)
         self.assertDictEqual(after['args'], args)
         self.assertEqual(after['return_value'], future)
+
+    def test_sample(self):
+        """Before/After solver sample events are dispatched with correct signatures."""
+
+        for solver in self.solvers:
+            with self.subTest("solver=%r" % solver.id):
+                self.subtest_sample(solver)
+
+
+class TestEventDispatchDecorator(unittest.TestCase):
+
+    def setUp(self):
+        # reset all event handlers
+        from dwave.cloud.events import _client_event_hooks_registry as reg
+        reg.update({k: [] for k in reg})
+
+    def test_decorator(self):
+        """Decorator adds on-entry and on-exit event calls, with correct args."""
+
+        class MockSampler:
+            @dispatches_events('sample')
+            def mock_sample(self, h, J, offset=0, fail=False, **kwargs):
+                if fail:
+                    raise ValueError
+                return offset + 1
+
+        mock_object = MockSampler()
+        h = [1, 1]
+        J = {(0, 1): 1}
+
+        def before(name, obj, args):
+            self.assertEqual(obj, mock_object)
+            args.pop('fail')
+            self.assertEqual(args, dict(h=h, J=J, offset=0, kwargs={}))
+
+        def after(name, obj, args, return_value=None, exception=None):
+            self.assertEqual(obj, mock_object)
+            fail = args.pop('fail')
+            self.assertEqual(args, dict(h=h, J=J, offset=0, kwargs={}))
+            if fail:
+                self.assertIsInstance(exception, ValueError)
+            else:
+                self.assertEqual(return_value, 1)
+
+        add_handler('before_sample', before)
+        add_handler('before_sample', after)
+
+        mock_object.mock_sample(h, J)
+        with self.assertRaises(ValueError):
+            mock_object.mock_sample(h, J, fail=True)