Merge branch 'bubbleboy14-master'

1 files changed, 59 insertions, 36 deletions

diff --git a/websocket/_app.py b/websocket/_app.py
index cd17a18..8a37be5 100644
--- a/websocket/_app.py
+++ b/websocket/_app.py
@@ -141,7 +141,7 @@ class WebSocketApp(object):
                     http_proxy_host=None, http_proxy_port=None,
                     http_no_proxy=None, http_proxy_auth=None,
                     skip_utf8_validation=False,
-                    host=None, origin=None):
+                    host=None, origin=None, dispatcher=None):
         """
         run event loop for WebSocket framework.
         This loop is infinite loop and is alive during websocket is available.
@@ -175,6 +175,22 @@ class WebSocketApp(object):
         close_frame = None
         self.keep_running = True
 
+        if not dispatcher:
+            dispatcher = self.create_dispatcher(ping_timeout)
+
+        def teardown():
+            if not self.keep_running:
+                return
+            if thread and thread.isAlive():
+                event.set()
+                thread.join()
+            self.keep_running = False
+            self.sock.close()
+            close_args = self._get_close_args(
+                close_frame.data if close_frame else None)
+            self._callback(self.on_close, *close_args)
+            self.sock = None
+
         try:
             self.sock = WebSocket(
                 self.get_mask_key, sockopt=sockopt, sslopt=sslopt,
@@ -196,53 +212,60 @@ class WebSocketApp(object):
                 thread.setDaemon(True)
                 thread.start()
 
-            while self.sock.connected:
-                r, w, e = select.select(
-                    (self.sock.sock, ), (), (), ping_timeout or 10) # Use a 10 second timeout to avoid to wait forever on close
+            def read():
                 if not self.keep_running:
-                    break
+                    return teardown()
 
-                if r:
-                    op_code, frame = self.sock.recv_data_frame(True)
-                    if op_code == ABNF.OPCODE_CLOSE:
-                        close_frame = frame
-                        break
-                    elif op_code == ABNF.OPCODE_PING:
-                        self._callback(self.on_ping, frame.data)
-                    elif op_code == ABNF.OPCODE_PONG:
-                        self.last_pong_tm = time.time()
-                        self._callback(self.on_pong, frame.data)
-                    elif op_code == ABNF.OPCODE_CONT and self.on_cont_message:
-                        self._callback(self.on_data, frame.data,
-                                       frame.opcode, frame.fin)
-                        self._callback(self.on_cont_message,
-                                       frame.data, frame.fin)
-                    else:
-                        data = frame.data
-                        if six.PY3 and op_code == ABNF.OPCODE_TEXT:
-                            data = data.decode("utf-8")
-                        self._callback(self.on_data, data, frame.opcode, True)
-                        self._callback(self.on_message, data)
+                op_code, frame = self.sock.recv_data_frame(True)
+                if op_code == ABNF.OPCODE_CLOSE:
+                    close_frame = frame
+                    return teardown()
+                elif op_code == ABNF.OPCODE_PING:
+                    self._callback(self.on_ping, frame.data)
+                elif op_code == ABNF.OPCODE_PONG:
+                    self.last_pong_tm = time.time()
+                    self._callback(self.on_pong, frame.data)
+                elif op_code == ABNF.OPCODE_CONT and self.on_cont_message:
+                    self._callback(self.on_data, frame.data,
+                                   frame.opcode, frame.fin)
+                    self._callback(self.on_cont_message,
+                                   frame.data, frame.fin)
+                else:
+                    data = frame.data
+                    if six.PY3 and op_code == ABNF.OPCODE_TEXT:
+                        data = data.decode("utf-8")
+                    self._callback(self.on_data, data, frame.opcode, True)
+                    self._callback(self.on_message, data)
 
                 if ping_timeout and self.last_ping_tm \
                         and time.time() - self.last_ping_tm > ping_timeout \
                         and self.last_ping_tm - self.last_pong_tm > ping_timeout:
                     raise WebSocketTimeoutException("ping/pong timed out")
+                return True
+
+            dispatcher.read(self.sock.sock, read)
+            dispatcher.dispatch()
         except (Exception, KeyboardInterrupt, SystemExit) as e:
             self._callback(self.on_error, e)
             if isinstance(e, SystemExit):
                 # propagate SystemExit further
                 raise
-        finally:
-            if thread and thread.isAlive():
-                event.set()
-                thread.join()
-            self.keep_running = False
-            self.sock.close()
-            close_args = self._get_close_args(
-                close_frame.data if close_frame else None)
-            self._callback(self.on_close, *close_args)
-            self.sock = None
+            teardown()
+
+    def create_dispatcher(self, ping_timeout):
+        class Dispatcher:
+            def read(_, sock, callback):
+                while self.sock.connected:
+                    r, w, e = select.select(
+                    (self.sock.sock, ), (), (), ping_timeout or 10) # Use a 10 second timeout to avoid to wait forever on close
+                    if r:
+                        callback()
+
+            def dispatch(_):
+                # do nothing
+                pass
+
+        return Dispatcher()
 
     def _get_close_args(self, data):
         """ this functions extracts the code, reason from the close body