Websockets

.github/workflows/lint.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ['3.7', '3.8', '3.9', '3.10']
+        python-version: ['3.8', '3.9', '3.10']
     name: Lint ${{ matrix.python-version }}
     steps:
       - uses: actions/checkout@v3

.github/workflows/release.yml

@@ -23,8 +23,11 @@ jobs:
         run: poetry config pypi-token.pypi ${{ secrets.POETRY_HTTP_BASIC_PYPI_PASSWORD }}
       - name: Install pypi deps
         run: poetry install
+      - name: Get tag
+        id: tag
+        uses: dawidd6/action-get-tag@v1
       - name: Version according to tag
-        run: poetry version $(git describe --tags)
+        run: poetry version ${{ steps.tag.outputs.tag }}
       - name: Build
         run: poetry build
       - name: Publish to PyPi

.github/workflows/test.yml

@@ -14,8 +14,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ['3.7', '3.8', '3.9', '3.10']
-    name: Lint ${{ matrix.python-version }}
+        python-version: ['3.8', '3.9', '3.10']
+    name: Unit test ${{ matrix.python-version }}
     steps:
       - uses: actions/checkout@v3
       - name: Setup Python

Makefile

@@ -6,7 +6,7 @@ YELLOW := $(shell tput -Txterm setaf 3)
 WHITE  := $(shell tput -Txterm setaf 7)
 RESET  := $(shell tput -Txterm sgr0)
 
-.PHONY: help lint style static test
+.PHONY: help lint style static test test_rest test_websocket
 
 ## Show help
 help:
@@ -22,16 +22,21 @@ help:
 
 ## Check code style
 style:
-	poetry run black $(if $(CI),--check,) polygon tests
+	poetry run black $(if $(CI),--check,) polygon test_*
 
 ## Check static types
 static:
-	poetry run mypy polygon tests
+	poetry run mypy polygon test_*
 
 ## Check code style and static types
 lint: style static
 
 ## Run unit tests
-test:
-	poetry run python -m unittest discover -s tests
+test_rest:
+	poetry run python -m unittest discover -s test_rest
+
+test_websocket:
+	poetry run python -m unittest discover -s test_websocket
+
+test: test_rest test_websocket
 

README.md

@@ -41,3 +41,20 @@ from polygon import RESTClient
 client = RESTClient() # Uses POLYGON_API_KEY env var. Can optionally supply your key.
 response = client.get_aggs("AAPL", 1, "day", "2005-04-01", "2005-04-04", raw=True)
 ```
+
+### Streaming websockets
+
+```python
+from polygon import WebSocketClient
+from polygon.websocket.models import Market, Feed, WebSocketMessage
+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')
+
+def handle_msg(msg: WebSocketMessage):
+	print(msg)
+
+asyncio.run(client.connect(handle_msg))
+```
+

docs/source/Getting-Started.rst

@@ -173,4 +173,18 @@ If it is a paginated :code:`list_*` response it's up to you to handle the "next_
 Websocket client usage
 ----------------------
 
-Coming soon.
+.. code-block:: python
+
+  from polygon import WebSocketClient
+  from polygon.websocket.models import Market, Feed, 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')
+
+  async def handle_msg(msg: List[WebSocketMessage]):
+    print(msg)
+
+  asyncio.run(client.connect(handle_msg))
+

docs/source/WebSockets.rst

@@ -0,0 +1,30 @@
+.. _websockets_header:
+
+WebSockets
+==========
+
+===========
+Init client
+===========
+.. automethod:: polygon.WebSocketClient.__init__
+
+============================
+Connect
+============================
+.. automethod:: polygon.WebSocketClient.connect
+
+============================
+Subscribe
+============================
+.. automethod:: polygon.WebSocketClient.subscribe
+
+============================
+Unsubscribe
+============================
+.. automethod:: polygon.WebSocketClient.unsubscribe
+
+============================
+Close
+============================
+.. automethod:: polygon.WebSocketClient.close
+

docs/source/index.rst

@@ -9,6 +9,7 @@ This documentation is for the Python client only. For details about the response
 
   Getting-Started
   Aggs
+  WebSockets
   Snapshot
   Quotes
   Reference

polygon/__init__.py

@@ -1 +1,2 @@
 from .rest import RESTClient
+from .websocket import WebSocketClient

polygon/rest/__init__.py

@@ -11,6 +11,12 @@
     ExchangesClient,
 )
 from .vx import VXClient
+from typing import Optional
+import os
+
+
+base = "https://api.polygon.io"
+env_key = "POLYGON_API_KEY"
 
 
 class RESTClient(
@@ -25,6 +31,31 @@ class RESTClient(
     ConditionsClient,
     ExchangesClient,
 ):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.vx = VXClient(*args, **kwargs)
+    def __init__(
+        self,
+        api_key: Optional[str] = os.getenv(env_key),
+        connect_timeout: float = 10.0,
+        read_timeout: float = 10.0,
+        num_pools: int = 10,
+        retries: int = 3,
+        base: str = base,
+        verbose: bool = False,
+    ):
+        super().__init__(
+            api_key=api_key,
+            connect_timeout=connect_timeout,
+            read_timeout=read_timeout,
+            num_pools=num_pools,
+            retries=retries,
+            base=base,
+            verbose=verbose,
+        )
+        self.vx = VXClient(
+            api_key=api_key,
+            connect_timeout=connect_timeout,
+            read_timeout=read_timeout,
+            num_pools=num_pools,
+            retries=retries,
+            base=base,
+            verbose=verbose,
+        )

polygon/rest/base.py

@@ -6,24 +6,21 @@
 from typing import Optional, Any, Dict
 from datetime import datetime
 
-base = "https://api.polygon.io"
-env_key = "POLYGON_API_KEY"
-
 
 class BaseClient:
     def __init__(
         self,
-        api_key: Optional[str] = os.getenv(env_key),
-        connect_timeout: float = 10.0,
-        read_timeout: float = 10.0,
-        num_pools: int = 10,
-        retries=3,
-        base: str = base,
-        verbose: bool = False,
+        api_key: Optional[str],
+        connect_timeout: float,
+        read_timeout: float,
+        num_pools: int,
+        retries: int,
+        base: str,
+        verbose: bool,
     ):
         if api_key is None:
             raise Exception(
-                f"Must specify env var {env_key} or pass api_key in constructor"
+                f"Must specify env var POLYGON_API_KEY or pass api_key in constructor"
             )
         self.API_KEY = api_key
         self.BASE = base

polygon/rest/quotes.py

@@ -1,5 +1,5 @@
 from .base import BaseClient
-from typing import Optional, Any, Dict, List, Union
+from typing import Optional, Any, Dict, List, Union, Iterator
 from .models import Quote, LastQuote, Sort, Order
 from urllib3 import HTTPResponse
 from datetime import datetime, date
@@ -19,7 +19,7 @@ def list_quotes(
         order: Optional[Union[str, Order]] = None,
         params: Optional[Dict[str, Any]] = None,
         raw: bool = False,
-    ) -> Union[List[Quote], HTTPResponse]:
+    ) -> Union[Iterator[Quote], HTTPResponse]:
         """
         Get quotes for a ticker symbol in a given time range.