make websockets have 2 apis (#163)

* make websockets have 2 apis

* lint

Author: clickingbuttons

README.md

@@ -5,13 +5,13 @@
 
 Python client for the [Polygon.io API](https://polygon.io).
 
-## Getting started
-### Install
+## Install
 
 `pip install polygon-api-client`
 
 Requires python version >= 3.7
 
+## REST getting started
 ### Getting aggs
 ```python
 from polygon import RESTClient
@@ -42,19 +42,86 @@ client = RESTClient() # Uses POLYGON_API_KEY env var. Can optionally supply your
 response = client.get_aggs("AAPL", 1, "day", "2005-04-01", "2005-04-04", raw=True)
 ```
 
-### Streaming websockets
+## WebSocket getting started
 
+### Simple synchronous callback
 ```python
 from polygon import WebSocketClient
-from polygon.websocket.models import Market, Feed, WebSocketMessage
-import asyncio
+from polygon.websocket.models import WebSocketMessage
+from typing import List
 
-client = WebSocketClient(market=Market.Stocks, feed=Feed.RealTime) # Uses POLYGON_API_KEY env var. Can optionally supply your key.
-client.subscribe('T.AAPL')
+c = WebSocketClient(subscriptions=['T.AAPL']) # Uses POLYGON_API_KEY env var. Can optionally supply your key.
 
-def handle_msg(msg: WebSocketMessage):
-	print(msg)
+def handle_msg(msgs: List[WebSocketMessage]):
+    for m in msgs:
+        print(m)
 
-asyncio.run(client.connect(handle_msg))
+c.run(handle_msg)
 ```
 
+### Synchronous aggregates
+```python
+from polygon import WebSocketClient
+from polygon.websocket.models import WebSocketMessage
+from typing import List
+
+class MessageHandler:
+    count = 0
+
+    def handle_msg(self, msgs: List[WebSocketMessage]):
+        for m in msgs:
+            if type(m) == EquityTrade:
+                print(self.count, m)
+                self.count += 1
+
+h = MessageHandler()
+
+def handle_msg(msgs: List[WebSocketMessage]):
+    h.handle_msg(msgs)
+
+c.run(handle_msg)
+```
+
+### Asynchronous callback
+```python
+from polygon import WebSocketClient
+from polygon.websocket.models import WebSocketMessage
+from typing import List
+
+c = WebSocketClient(subscriptions=['T.AAPL']) # Uses POLYGON_API_KEY env var. Can optionally supply your key.
+
+async def handle_msg(msgs: List[WebSocketMessage]):
+    for m in msgs:
+        print(m)
+
+async def timeout():
+    await asyncio.sleep(1)
+    print('unsubscribe_all')
+    c.unsubscribe_all()
+    await asyncio.sleep(1)
+    print('close')
+    await c.close()
+
+async def main():
+    await asyncio.gather(
+        c.connect(handle_msg),
+        timeout()
+    )
+
+asyncio.run(main())
+```
+
+### Getting raw response
+```python
+from polygon import WebSocketClient
+from polygon.websocket.models import WebSocketMessage
+from typing import Union
+import json
+
+c = WebSocketClient(subscriptions=['T.*'], raw=True)
+
+def handle_msg(msgs: Union[str, bytes]):
+		print(json.loads(msgs))
+
+c.run(handle_msg)
+```

docs/source/Getting-Started.rst

@@ -175,23 +175,100 @@ If it is a paginated :code:`list_*` response it's up to you to handle the "next_
   # }'
 
 
-Websocket client usage
+WebSocket client usage
 ----------------------
 
+.. automethod:: polygon.WebSocketClient.__init__
+
+The simplest way to use the websocket client is to just provide a callback:
+
 .. code-block:: python
 
   from polygon import WebSocketClient
-  from polygon.websocket.models import Market, Feed, WebSocketMessage
+  from polygon.websocket.models import WebSocketMessage
   from typing import List
-  import asyncio
 
-  client = WebSocketClient(market=Market.Stocks, feed=Feed.RealTime) # Uses POLYGON_API_KEY env var. Can optionally supply your key.
-  client.subscribe('T.AAPL')
+  c = WebSocketClient(subscriptions=['T.AAPL'])
 
-  async def handle_msg(msg: List[WebSocketMessage]):
-    print(msg)
+  def handle_msg(msgs: List[WebSocketMessage]):
+      for m in msgs:
+          print(m)
 
-  asyncio.run(client.connect(handle_msg))
+  c.run(handle_msg)
 
 .. note::
-  Raises :code:`AuthError` if invalid API key is provided.
+  Raises :code:`AuthError` if invalid API key is provided.
+
+If you want to capture state you can use a global variable inside the callback.
+Alternatively, you can wrap a class method in a closure.
+
+.. code-block:: python
+
+  from polygon import WebSocketClient
+  from polygon.websocket.models import WebSocketMessage
+  from typing import List
+
+  class MessageHandler:
+      count = 0
+
+      def handle_msg(self, msgs: List[WebSocketMessage]):
+          for m in msgs:
+              if type(m) == EquityTrade:
+                  print(self.count, m)
+                  self.count += 1
+
+  h = MessageHandler()
+
+  def handle_msg(msgs: List[WebSocketMessage]):
+      h.handle_msg(msgs)
+
+  c.run(handle_msg)
+
+Under the hood our client uses an asynchronous runtime. To manage the runtime
+yourself (including unsubscribing and subscribing) you'll need to use asyncio
+and the :code:`.connect` method:
+
+.. code-block:: python
+
+  from polygon import WebSocketClient
+  from polygon.websocket.models import WebSocketMessage
+  from typing import List
+
+  c = WebSocketClient(subscriptions=['T.AAPL']) # Uses POLYGON_API_KEY env var. Can optionally supply your key.
+
+  async def handle_msg(msgs: List[WebSocketMessage]):
+      for m in msgs:
+          print(m)
+
+  async def timeout():
+      await asyncio.sleep(1)
+      print('unsubscribe_all')
+      c.unsubscribe_all()
+      await asyncio.sleep(1)
+      print('close')
+      await c.close()
+
+  async def main():
+      await asyncio.gather(
+          c.connect(handle_msg),
+          timeout()
+      )
+
+  asyncio.run(main())
+
+To handle raw messages yourself pass `raw=True`:
+
+.. code-block:: python
+
+  from polygon import WebSocketClient
+  from polygon.websocket.models import WebSocketMessage
+  from typing import Union
+  import json
+
+  c = WebSocketClient(subscriptions=['T.*'], raw=True)
+
+  def handle_msg(msgs: Union[str, bytes]):
+      print(json.loads(msgs))
+
+  c.run(handle_msg)
+

docs/source/WebSocket.rst

@@ -13,6 +13,11 @@ Connect
 ============================
 .. automethod:: polygon.WebSocketClient.connect
 
+============================
+Run
+============================
+.. automethod:: polygon.WebSocketClient.run
+
 ============================
 Subscribe
 ============================

polygon/websocket/__init__.py

@@ -2,13 +2,12 @@
 from enum import Enum
 from typing import Optional, Union, List, Set, Callable, Awaitable
 import json
-import inspect
+import asyncio
 import ssl
 import certifi
 from .models import *
 from websockets.client import connect, WebSocketClientProtocol
 from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError
-from websockets.typing import Data
 
 env_key = "POLYGON_API_KEY"
 
@@ -66,7 +65,10 @@ def __init__(
     # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection
     async def connect(
         self,
-        processor: Callable[[Union[List[WebSocketMessage], Data]], Optional[Awaitable]],
+        processor: Union[
+            Callable[[List[WebSocketMessage]], Awaitable],
+            Callable[[Union[str, bytes]], Awaitable],
+        ],
         close_timeout: int = 1,
         **kwargs,
     ):
@@ -78,7 +80,6 @@ async def connect(
         :raises AuthError: If invalid API key is supplied.
         """
         reconnects = 0
-        isasync = inspect.iscoroutinefunction(processor)
         if self.verbose:
             print("connect:", self.url)
         # darwin needs some extra <3
@@ -116,7 +117,9 @@ async def connect(
                         self.subs = set(self.scheduled_subs)
                         self.schedule_resub = False
 
-                    cmsg: Union[List[WebSocketMessage], Data] = await s.recv()
+                    cmsg: Union[
+                        List[WebSocketMessage], Union[str, bytes]
+                    ] = await s.recv()
                     # we know cmsg is Data
                     msgJson = json.loads(cmsg)  # type: ignore
                     for m in msgJson:
@@ -128,11 +131,7 @@ async def connect(
                         cmsg = parse(msgJson)
 
                     if len(cmsg) > 0:
-                        if isasync:
-                            # we know processor is Awaitable
-                            await processor(cmsg)  # type: ignore
-                        else:
-                            processor(cmsg)
+                        await processor(cmsg)  # type: ignore
             except ConnectionClosedOK:
                 if self.verbose:
                     print("connection closed (OK)")
@@ -145,6 +144,28 @@ async def connect(
                     return
                 continue
 
+    def run(
+        self,
+        handle_msg: Union[
+            Callable[[List[WebSocketMessage]], None],
+            Callable[[Union[str, bytes]], None],
+        ],
+        close_timeout: int = 1,
+        **kwargs,
+    ):
+        """
+        Connect to websocket server and run `processor(msg)` on every new `msg`. Synchronous version of `.connect`.
+
+        :param processor: The callback to process messages.
+        :param close_timeout: How long to wait for handshake when calling .close.
+        :raises AuthError: If invalid API key is supplied.
+        """
+
+        async def handle_msg_wrapper(msgs):
+            handle_msg(msgs)
+
+        asyncio.run(self.connect(handle_msg_wrapper, close_timeout, **kwargs))
+
     async def _subscribe(self, topics: Union[List[str], Set[str]]):
         if self.websocket is None or len(topics) == 0:
             return