Merge branch 'master' of https://github.com/danpaquin/gdax-python

Author: lirenz

README.md

@@ -240,16 +240,19 @@ wsClient.close()
 ### WebsocketClient + Mongodb
 The ```WebsocketClient``` now supports data gathering via MongoDB. Given a
 MongoDB collection, the ```WebsocketClient``` will stream results directly into
-the database collection. 
+the database collection.
 ```python
 # import PyMongo and connect to a local, running Mongo instance
 from pymongo import MongoClient
+import gdax
 mongo_client = MongoClient('mongodb://localhost:27017/')
+
 # specify the database and collection
 db = mongo_client.cryptocurrency_database
 BTC_collection = db.BTC_collection
+
 # instantiate a WebsocketClient instance, with a Mongo collection as a parameter
-wsClient = WebsocketClient(url="wss://ws-feed.gdax.com", products="BTC-USD",
+wsClient = gdax.WebsocketClient(url="wss://ws-feed.gdax.com", products="BTC-USD",
     mongo_collection=BTC_collection, should_print=False)
 wsClient.start()
 ```
@@ -261,9 +264,9 @@ can react to the data streaming in.  The current client is a template used for
 illustration purposes only.
 
 
-- onOpen - called once, *immediately before* the socket connection is made, this 
+- onOpen - called once, *immediately before* the socket connection is made, this
 is where you want to add initial parameters.
-- onMessage - called once for every message that arrives and accepts one 
+- onMessage - called once for every message that arrives and accepts one
 argument that contains the message of dict type.
 - onClose - called once after the websocket has been closed.
 - close - call this method to close the websocket connection (do not overwrite).

contributors.txt

@@ -3,3 +3,4 @@ Leonard Lin
 Jeff Gibson
 David Caseria
 Paul Mestemaker
+Drew Rice

gdax/order_book.py

@@ -36,38 +36,39 @@ def on_open(self):
     def on_close(self):
         print("\n-- OrderBook Socket Closed! --")
 
+    def reset_book(self):
+        self._asks = RBTree()
+        self._bids = RBTree()
+        res = self._client.get_product_order_book(product_id=self.product_id, level=3)
+        for bid in res['bids']:
+            self.add({
+                'id': bid[2],
+                'side': 'buy',
+                'price': Decimal(bid[0]),
+                'size': Decimal(bid[1])
+            })
+        for ask in res['asks']:
+            self.add({
+                'id': ask[2],
+                'side': 'sell',
+                'price': Decimal(ask[0]),
+                'size': Decimal(ask[1])
+            })
+        self._sequence = res['sequence']
+
     def on_message(self, message):
         if self._log_to:
             pickle.dump(message, self._log_to)
 
         sequence = message['sequence']
         if self._sequence == -1:
-            self._asks = RBTree()
-            self._bids = RBTree()
-            res = self._client.get_product_order_book(product_id=self.product_id, level=3)
-            for bid in res['bids']:
-                self.add({
-                    'id': bid[2],
-                    'side': 'buy',
-                    'price': Decimal(bid[0]),
-                    'size': Decimal(bid[1])
-                })
-            for ask in res['asks']:
-                self.add({
-                    'id': ask[2],
-                    'side': 'sell',
-                    'price': Decimal(ask[0]),
-                    'size': Decimal(ask[1])
-                })
-            self._sequence = res['sequence']
-
+            self.reset_book()
+            return
         if sequence <= self._sequence:
             # ignore older messages (e.g. before order book initialization from getProductOrderBook)
             return
         elif sequence > self._sequence + 1:
-            print('Error: messages missing ({} - {}). Re-initializing websocket.'.format(sequence, self._sequence))
-            self.close()
-            self.start()
+            self.on_sequence_gap(self._sequence, sequence)
             return
 
         msg_type = message['type']
@@ -83,18 +84,11 @@ def on_message(self, message):
 
         self._sequence = sequence
 
-        # bid = self.get_bid()
-        # bids = self.get_bids(bid)
-        # bid_depth = sum([b['size'] for b in bids])
-        # ask = self.get_ask()
-        # asks = self.get_asks(ask)
-        # ask_depth = sum([a['size'] for a in asks])
-        # print('bid: %f @ %f - ask: %f @ %f' % (bid_depth, bid, ask_depth, ask))
+    def on_sequence_gap(self, gap_start, gap_end):
+        self.reset_book()
+        print('Error: messages missing ({} - {}). Re-initializing  book at sequence.'.format(
+            gap_start, gap_end, self._sequence))
 
-    def on_error(self, e):
-        self._sequence = -1
-        self.close()
-        self.start()
 
     def add(self, order):
         order = {
@@ -250,6 +244,7 @@ def set_bids(self, price, bids):
 
 
 if __name__ == '__main__':
+    import sys
     import time
     import datetime as dt
 
@@ -286,10 +281,18 @@ def on_message(self, message):
                 self._ask = ask
                 self._bid_depth = bid_depth
                 self._ask_depth = ask_depth
-                print('{}\tbid: {:.3f} @ {:.2f}\task: {:.3f} @ {:.2f}'.format(dt.datetime.now(), bid_depth, bid,
-                                                                              ask_depth, ask))
+                print('{} {} bid: {:.3f} @ {:.2f}\task: {:.3f} @ {:.2f}'.format(
+                    dt.datetime.now(), self.product_id, bid_depth, bid, ask_depth, ask))
 
     order_book = OrderBookConsole()
     order_book.start()
-    time.sleep(10)
-    order_book.close()
+    try:
+        while True:
+            time.sleep(10)
+    except KeyboardInterrupt:
+        order_book.close()
+
+    if order_book.error:
+        sys.exit(1)
+    else:
+        sys.exit(0)

gdax/websocket_client.py

@@ -16,13 +16,15 @@
 from pymongo import MongoClient
 
 class WebsocketClient(object):
-    def __init__(self, url="wss://ws-feed.gdax.com", products=None, message_type="subscribe",
-            mongo_collection=None, should_print=True, auth=False, api_key="", api_secret="", api_passphrase="", channels=None):
+    def __init__(self, url="wss://ws-feed.gdax.com", products=None, message_type="subscribe", mongo_collection=None, 
+        should_print=True, auth=False, api_key="", api_secret="", api_passphrase="", channels=None):
+
         self.url = url
         self.products = products
         self.channels = channels
         self.type = message_type
         self.stop = False
+        self.error = None
         self.ws = None
         self.thread = None
         self.auth = auth
@@ -36,6 +38,7 @@ def start(self):
         def _go():
             self._connect()
             self._listen()
+            self._disconnect()
 
         self.stop = False
         self.on_open()
@@ -85,22 +88,29 @@ def _listen(self):
                 if int(time.time() % 30) == 0:
                     # Set a 30 second ping to keep connection alive
                     self.ws.ping("keepalive")
-                msg = json.loads(self.ws.recv())
+                data = self.ws.recv()
+                msg = json.loads(data)
+            except ValueError as e:
+                self.on_error(e)
             except Exception as e:
                 self.on_error(e)
             else:
                 self.on_message(msg)
 
+    def _disconnect(self):
+        if self.type == "heartbeat":
+            self.ws.send(json.dumps({"type": "heartbeat", "on": False}))
+        try:
+            if self.ws:
+                self.ws.close()
+        except WebSocketConnectionClosedException as e:
+            pass
+
+        self.on_close()
+
     def close(self):
-        if not self.stop:
-            self.on_close()
-            self.stop = True
-            self.thread.join()
-            try:
-                if self.ws:
-                    self.ws.close()
-            except WebSocketConnectionClosedException as e:
-                pass
+        self.stop = True
+        self.thread.join()
 
     def on_open(self):
         if self.should_print:
@@ -116,10 +126,13 @@ def on_message(self, msg):
         if self.mongo_collection: # dump JSON to given mongo collection
             self.mongo_collection.insert_one(msg)
 
-    def on_error(self, e):
-        print(e)
+    def on_error(self, e, data=None):
+        self.error = e
+        self.stop
+        print('{} - data: {}'.format(e, data))
 
 if __name__ == "__main__":
+    import sys
     import gdax
     import time
 
@@ -131,8 +144,7 @@ def on_open(self):
             print("Let's count the messages!")
 
         def on_message(self, msg):
-            if 'price' in msg and 'type' in msg:
-                print("Message type:", msg["type"], "\t@ %.3f" % float(msg["price"]))
+            print(json.dumps(msg, indent=4, sort_keys=True))
             self.message_count += 1
 
         def on_close(self):
@@ -141,9 +153,14 @@ def on_close(self):
     wsClient = MyWebsocketClient()
     wsClient.start()
     print(wsClient.url, wsClient.products)
-    # Do some logic with the data
-    while wsClient.message_count < 10000:
-        print("\nMessageCount =", "%i \n" % wsClient.message_count)
-        time.sleep(1)
-
-    wsClient.close()
+    try:
+        while True:
+            print("\nMessageCount =", "%i \n" % wsClient.message_count)
+            time.sleep(1)
+    except KeyboardInterrupt:
+        wsClient.close()
+
+    if wsClient.error:
+        sys.exit(1)
+    else:
+        sys.exit(0)

requirements.txt

@@ -2,4 +2,4 @@ bintrees==2.0.7
 requests==2.13.0
 six==1.10.0
 websocket-client==0.40.0
-pymongo
+pymongo==3.5.1

setup.py

@@ -7,6 +7,7 @@
     'requests==2.13.0',
     'six==1.10.0',
     'websocket-client==0.40.0',
+    'pymongo==3.5.1'
 ]
 
 tests_require = [