Implement an error callback and arrange for reconnection

inkblot · inkblot · commit ceb807c01ec2 · 2016-11-24T07:10:03.000-06:00
* Splits the `setup()` method into separate `connect()` and `listen()` methods.
* Adjust the `connect()` and `close()` methods so that the websocket may be reconnected and stay active.
* Adds an `on_error()` callback that is invoked in the `except` block of `listen()` with default behavior of closing the connection.
* Makes WebSocketClient a new-style class.

These changes together are intended to make it possible to reconnect a disconnected client using an `on_error()` override like this:

```
    def on_error(self, e):
        if isinstance(e, WebSocketConnectionClosedException):
            while not self.stop:
                try:
                    self.close()
                    self.connect()
                except Exception as e:
                    logger.error(repr(e))
                else:
                    break
        else:
            logger.error(repr(e))
```
diff --git a/GDAX/WebsocketClient.py b/GDAX/WebsocketClient.py
@@ -8,27 +8,31 @@
 from   threading import Thread
 from   websocket import create_connection
 
-class WebsocketClient():
+class WebsocketClient(object):
     def __init__(self, ws_url="wss://ws-feed-public.sandbox.gdax.com", product_id="BTC-USD"):
         if ws_url[-1] == "/":
             self.url = ws_url[:-1]
         else:
             self.url = ws_url
-        self.stop = False
         self.product_id = product_id
-        self.thread = Thread(target=self.setup)
+
+        def go():
+            self.connect()
+            self.listen()
+
+        self.thread = Thread(target=go)
         self.thread.start()
 
-    def setup(self):
+    def connect(self):
         self.open()
+        self.stop = False
         self.ws = create_connection(self.url)
         if type(self.product_id) is list:
             #product_ids - plural for multiple products
             subParams = json.dumps({"type": "subscribe", "product_ids": self.product_id})
         else:
             subParams = json.dumps({"type": "subscribe", "product_id": self.product_id})
         self.ws.send(subParams)
-        self.listen()
 
     def open(self):
         print("-- Subscribed! --")
@@ -38,16 +42,21 @@ def listen(self):
             try:
                 msg = json.loads(self.ws.recv())
             except Exception as e:
-                #print e
-                break
+                self.on_error(e)
             else:
                 self.message(msg)
 
+    def on_error(self, e):
+        self.close()
+
     def message(self, msg):
         print(msg)
 
     def close(self):
-        self.ws.close()
+        self.stop = True
+        if self.ws and self.ws.connected:
+            self.ws.close()
+        self.ws = None
         self.closed()
 
     def closed(self):
