Open pandora box: allow hooks for patterns

bombela · bombela · commit 049c0b3e83ef · 2014-09-10T01:51:37.000-07:00
diff --git a/zerorpc/context.py b/zerorpc/context.py
@@ -44,7 +44,8 @@ def __init__(self):
             'server_inspect_exception': [],
             'client_handle_remote_error': [],
             'client_before_request': [],
-            'client_after_request': []
+            'client_after_request': [],
+            'client_patterns_list': [],
         }
         self._reset_msgid()
 
@@ -217,3 +218,8 @@ def hook_client_after_request(self, request_event, reply_event, exception=None):
         """
         for functor in self._hooks['client_after_request']:
             functor(request_event, reply_event, exception)
+
+    def hook_client_patterns_list(self, patterns):
+        for functor in self._hooks['client_patterns_list']:
+            patterns = functor(patterns)
+        return patterns
diff --git a/zerorpc/core.py b/zerorpc/core.py
@@ -206,27 +206,31 @@ def _handle_remote_error(self, event):
         return exception
 
     def _select_pattern(self, event):
-        for pattern in patterns.patterns_list:
+        for pattern in self._context.hook_client_patterns_list(
+                patterns.patterns_list):
             if pattern.accept_answer(event):
                 return pattern
-        msg = 'Unable to find a pattern for: {0}'.format(event)
-        raise RuntimeError(msg)
+        return None
 
     def _process_response(self, request_event, bufchan, timeout):
-        try:
-            reply_event = bufchan.recv(timeout)
-            pattern = self._select_pattern(reply_event)
-            return pattern.process_answer(self._context, bufchan, request_event,
-                    reply_event, self._handle_remote_error)
-        except TimeoutExpired:
+        def raise_error(ex):
             bufchan.close()
-            ex = TimeoutExpired(timeout,
-                    'calling remote method {0}'.format(request_event.name))
             self._context.hook_client_after_request(request_event, None, ex)
             raise ex
-        except:
-            bufchan.close()
-            raise
+
+        try:
+            reply_event = bufchan.recv(timeout=timeout)
+        except TimeoutExpired:
+            raise_error(TimeoutExpired(timeout,
+                    'calling remote method {0}'.format(request_event.name)))
+
+        pattern = self._select_pattern(reply_event)
+        if pattern is None:
+            raise_error(RuntimeError(
+                'Unable to find a pattern for: {0}'.format(request_event)))
+
+        return pattern.process_answer(self._context, bufchan, request_event,
+                reply_event, self._handle_remote_error)
 
     def __call__(self, method, *args, **kargs):
         timeout = kargs.get('timeout', self._timeout)
