updated WebsocketClient, requirements.txt

gdax/websocket_client.py

@@ -12,10 +12,11 @@
 import time
 from threading import Thread
 from websocket import create_connection, WebSocketConnectionClosedException
-
+from pymongo import MongoClient
 
 class WebsocketClient(object):
-    def __init__(self, url="wss://ws-feed.gdax.com", products=None, message_type="subscribe", auth=False, api_key="", api_secret="", api_passphrase=""):
+    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=""):
         self.url = url
         self.products = products
         self.type = message_type
@@ -26,6 +27,7 @@ def __init__(self, url="wss://ws-feed.gdax.com", products=None, message_type="su
         self.api_key = api_key
         self.api_secret = api_secret
         self.api_passphrase = api_passphrase
+        self.mongo_collection = mongo_collection
 
     def start(self):
         def _go():
@@ -99,7 +101,10 @@ def on_close(self):
         print("\n-- Socket Closed --")
 
     def on_message(self, msg):
-        print(msg)
+        if should_print:
+            print(msg)
+        if self.mongo_collection: # dump JSON to given mongo collection
+            self.mongo_collection.insert_one(msg)
 
     def on_error(self, e):
         print(e)

requirements.txt

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