Merge pull request #272 from brianboonstra/order_book_trees

Order book trees

Author: danpaquin

gdax/order_book.py

@@ -4,7 +4,7 @@
 #
 # Live order book updated from the gdax Websocket Feed
 
-from bintrees import RBTree
+from sortedcontainers import SortedDict
 from decimal import Decimal
 import pickle
 
@@ -15,8 +15,8 @@
 class OrderBook(WebsocketClient):
     def __init__(self, product_id='BTC-USD', log_to=None):
         super(OrderBook, self).__init__(products=product_id)
-        self._asks = RBTree()
-        self._bids = RBTree()
+        self._asks = SortedDict()
+        self._bids = SortedDict()
         self._client = PublicClient()
         self._sequence = -1
         self._log_to = log_to
@@ -37,8 +37,8 @@ def on_close(self):
         print("\n-- OrderBook Socket Closed! --")
 
     def reset_book(self):
-        self._asks = RBTree()
-        self._bids = RBTree()
+        self._asks = SortedDict()
+        self._bids = SortedDict()
         res = self._client.get_product_order_book(product_id=self.product_id, level=3)
         for bid in res['bids']:
             self.add({
@@ -219,28 +219,28 @@ def get_current_book(self):
         return result
 
     def get_ask(self):
-        return self._asks.min_key()
+        return self._asks.peekitem(0)[0]
 
     def get_asks(self, price):
         return self._asks.get(price)
 
     def remove_asks(self, price):
-        self._asks.remove(price)
+        del self._asks[price]
 
     def set_asks(self, price, asks):
-        self._asks.insert(price, asks)
+        self._asks[price] = asks
 
     def get_bid(self):
-        return self._bids.max_key()
+        return self._bids.peekitem(-1)[0]
 
     def get_bids(self, price):
         return self._bids.get(price)
 
     def remove_bids(self, price):
-        self._bids.remove(price)
+        del self._bids[price]
 
     def set_bids(self, price, bids):
-        self._bids.insert(price, bids)
+        self._bids[price] = bids
 
 
 if __name__ == '__main__':

requirements.txt

@@ -1,4 +1,4 @@
-bintrees==2.0.7
+sortedcontainers>=1.5.9
 requests==2.13.0
 six==1.10.0
 websocket-client==0.40.0

tests/test_public_client.py

@@ -34,7 +34,7 @@ def test_get_product_order_book(self, client, level):
         if level is 2 and (len(r['asks']) > 50 or len(r['bids']) > 50):
             pytest.fail('Fail: Level 2 should only return the top 50 asks and bids')
 
-        if level is 2 and (len(r['asks']) < 50 or len(r['bids']) < 50):
+        if level is 3 and (len(r['asks']) < 50 or len(r['bids']) < 50):
             pytest.fail('Fail: Level 3 should return the full order book')
 
     def test_get_product_ticker(self, client):