From 1c3e02a70a9566f34148613005458da218069014 Mon Sep 17 00:00:00 2001
From: 0xfMissingNo <0xfmissingno@protonmail.com>
Date: Thu, 23 Sep 2021 12:46:47 -0700
Subject: [PATCH] prototyping threaded ws client

---
 dydx3/dydx_client.py       | 17 ++++++++++
 dydx3/modules/websocket.py | 63 ++++++++++++++++++++++++++++++++++++++
 requirements.txt           |  1 +
 3 files changed, 81 insertions(+)
 create mode 100644 dydx3/modules/websocket.py

diff --git a/dydx3/dydx_client.py b/dydx3/dydx_client.py
index 204146d..8b9720c 100644
--- a/dydx3/dydx_client.py
+++ b/dydx3/dydx_client.py
@@ -8,6 +8,7 @@
 from dydx3.modules.private import Private
 from dydx3.modules.public import Public
 from dydx3.modules.onboarding import Onboarding
+from dydx3.modules.websocket import WSClient
 from dydx3.starkex.helpers import private_key_to_public_key_pair_hex
 from dydx3.starkex.starkex_resources.cpp_signature import (
     get_cpp_lib,
@@ -78,6 +79,7 @@ def __init__(
         self._api_keys = None
         self._eth = None
         self._onboarding = None
+        self._ws = None
 
         # Derive the public keys.
         if stark_private_key is not None:
@@ -118,6 +120,21 @@ def __init__(
                     e,
                 )
 
+    @property
+    def ws(self):
+        '''
+        Get the ws client module, used for interacting with websocket endpoints.
+        '''
+        if self._ws:
+            return self._ws
+        host = self.host
+        if host.startswith('https'):
+            host = host.replace('https', 'wss')
+        if host.startswith('http'):
+            host = host.replace('http', 'wss')
+        self._ws = WSClient(host + '/v3/ws')
+        return self._ws
+
     @property
     def public(self):
         '''
diff --git a/dydx3/modules/websocket.py b/dydx3/modules/websocket.py
new file mode 100644
index 0000000..d06f9a8
--- /dev/null
+++ b/dydx3/modules/websocket.py
@@ -0,0 +1,63 @@
+import json
+from threading import Thread
+import websocket
+
+
+class WSClient(websocket.WebSocketApp):
+
+    callback = print
+
+    @staticmethod
+    def on_open(ws):
+        WSClient.callback('-- connection established --')
+
+    @staticmethod
+    def on_close(ws, close_status_code, close_msg):
+        WSClient.callback('-- connection closed --')
+
+    @staticmethod
+    def on_message(ws, message):
+        if json.loads(message)['type'] == 'channel_data':
+            WSClient.callback(message)
+
+    @staticmethod
+    def on_error(ws, message):
+        WSClient.callback(message)
+
+    def __init__(self, host):
+        self.uri = host
+        super().__init__(
+            self.uri,
+            on_open=self.on_open,
+            on_close=self.on_close,
+            on_message=self.on_message,
+            on_error=self.on_error
+            )
+        self._thread = None
+    
+    def subscribe_to_orderbook(self, trading_pair):
+        self.send({'type': "subscribe", "channel": "v3_orderbook", "id": trading_pair})
+
+    def subscribe_to_trades(self, trading_pair):
+        self.send({'type': "subscribe", "channel": "v3_trades", "id": trading_pair})
+
+    def subscribe_to_markets(self):
+        self.send({'type': "subscribe", "channel": "v3_markets"})
+    
+    def send(self, data, *args):
+        super().send(json.dumps(data), *args)
+
+    def start(self):
+        if self._thread:
+            return self._thread
+        self._thread = Thread(target=self.run_forever, args=())
+        self._thread.start()
+
+    def stop(self):
+        if not self._thread:
+            return
+        try:
+            self.close()
+            self._thread.join()
+        finally:
+            self._thread = None
diff --git a/requirements.txt b/requirements.txt
index ff19dee..db4059c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -11,3 +11,4 @@ six==1.12
 sympy==1.6
 tox==3.13.2
 web3>=5.0.0,<6.0.0
+websocket-client==1.2.1
\ No newline at end of file