Merge pull request #1 from Polygon-io/websocket-mvp

Initial take on the websocket client

Author: qrpike

README.md

@@ -1 +1,36 @@
 # A Python client library for Polgyon's WebSocket and RESTful APIs
+
+Currently this repo only supports the WebSocket API
+
+## Getting Started
+
+For a basic product overview, check out our [setup and use documentation](https://polygon.io/sockets)
+
+
+## Simple Demo
+```python
+import time
+
+
+from polygon_client import WebSocketClient, STOCKS_CLUSTER
+
+
+def my_customer_process_message(message):
+    print("this is my custom message processing", message)
+
+
+def main():
+    key = 'your api key'
+    my_client = WebSocketClient(STOCKS_CLUSTER, key, my_customer_process_message)
+    my_client.run_async()
+
+    my_client.subscribe("T.MSFT", "T.AAPL", "T.AMD", "T.NVDA")
+    time.sleep(2)
+
+    my_client.close_connection()
+
+
+if __name__ == "__main__":
+    main()
+
+```

example.py

@@ -0,0 +1,22 @@
+import time
+
+from polygon_client import WebSocketClient, STOCKS_CLUSTER
+
+
+def my_customer_process_message(message):
+    print("this is my custom message processing", message)
+
+
+def main():
+    key = 'your api key'
+    my_client = WebSocketClient(STOCKS_CLUSTER, key, my_customer_process_message)
+    my_client.run_async()
+
+    my_client.subscribe("T.MSFT", "T.AAPL", "T.AMD", "T.NVDA")
+    time.sleep(1)
+
+    my_client.close_connection()
+
+
+if __name__ == "__main__":
+    main()

polygon_client/__init__.py

@@ -0,0 +1 @@
+from .websocket_client import WebSocketClient, STOCKS_CLUSTER, FOREX_CLUSTER, CRYPTO_CLUSTER

polygon_client/websocket_client.py

@@ -0,0 +1,102 @@
+import signal
+import threading
+from typing import Optional, Callable
+
+import websocket
+
+STOCKS_CLUSTER = "stocks"
+FOREX_CLUSTER = "forex"
+CRYPTO_CLUSTER = "crypto"
+
+
+class WebSocketClient(object):
+    DEFAULT_HOST = 'socket.polygon.io'
+
+    # TODO: Either an instance of the client couples 1:1 with the cluster or an instance of the Client couples 1:3 with
+    #  the 3 possible clusters (I think I like client per, but then a problem is the user can make multiple clients for
+    #  the same cluster and that's not desirable behavior,
+    #  somehow keeping track with multiple Client instances will be the difficulty)
+    def __init__(self, cluster: str, auth_key: str, process_message: Optional[Callable[[str], None]] = None):
+        self._host = self.DEFAULT_HOST
+        self.url = f"wss://{self._host}/{cluster}"
+        self.ws: websocket.WebSocketApp = websocket.WebSocketApp(self.url, on_open=self._default_on_open(),
+                                                                 on_close=self._default_on_close,
+                                                                 on_error=self._default_on_error,
+                                                                 on_message=self._default_on_message())
+        self.auth_key = auth_key
+
+        self.process_message = process_message
+
+        # being authenticated is an event that must occur before any other action is sent to the server
+        self._authenticated = threading.Event()
+        # self._run_thread is only set if the client is run asynchronously
+        self._run_thread: Optional[threading.Thread] = None
+
+        signal.signal(signal.SIGINT, self._cleanup_signal_handler())
+        signal.signal(signal.SIGTERM, self._cleanup_signal_handler())
+
+    def run(self):
+        self.ws.run_forever()
+
+    def run_async(self):
+        self._run_thread = threading.Thread(target=self.run)
+        self._run_thread.start()
+
+    def close_connection(self):
+        self.ws.close()
+        if self._run_thread:
+            self._run_thread.join()
+
+    def subscribe(self, *params):
+        # TODO: make this a decorator or context manager
+        self._authenticated.wait()
+
+        sub_message = '{"action":"subscribe","params":"%s"}' % self._format_params(params)
+        self.ws.send(sub_message)
+
+    def unsubscribe(self, *params):
+        # TODO: make this a decorator or context manager
+        self._authenticated.wait()
+
+        sub_message = '{"action":"unsubscribe","params":"%s"}' % self._format_params(params)
+        self.ws.send(sub_message)
+
+    def _cleanup_signal_handler(self):
+        return lambda signalnum, frame: self.close_connection()
+
+    def _authenticate(self, ws):
+        ws.send('{"action":"auth","params":"%s"}' % self.auth_key)
+        self._authenticated.set()
+
+    @staticmethod
+    def _format_params(params):
+        return ",".join(params)
+
+    @property
+    def process_message(self):
+        return self.__process_message
+
+    @process_message.setter
+    def process_message(self, pm):
+        if pm:
+            self.__process_message = pm
+            self.ws.on_message = lambda ws, message: self.__process_message(message)
+
+    def _default_on_message(self):
+        return lambda ws, message: self._default_process_message(message)
+
+    @staticmethod
+    def _default_process_message(message):
+        print(message)
+
+    def _default_on_open(self):
+        def f(ws):
+            self._authenticate(ws)
+
+    @staticmethod
+    def _default_on_error(ws, error):
+        print("error:", error)
+
+    @staticmethod
+    def _default_on_close(ws):
+        print("### closed ###")

requirements.txt

@@ -0,0 +1,12 @@
+aiofiles==0.4.0
+aiohttp==3.6.2
+async-timeout==3.0.1
+attrs==19.3.0
+certifi==2019.9.11
+chardet==3.0.4
+idna==2.8
+multidict==4.5.2
+six==1.12.0
+websocket-client==0.56.0
+websockets==8.0.2
+yarl==1.3.0