allow datetime or int for aggs, lint tests (#123)

* allow datetime or int for aggs, lint tests

* run --check in CI

Author: clickingbuttons

Makefile

@@ -22,11 +22,11 @@ help:
 
 ## Check code style
 style:
-	poetry run black polygon
+	poetry run black $(if $(CI),--check,) polygon tests
 
 ## Check static types
 static:
-	poetry run mypy polygon
+	poetry run mypy polygon tests
 
 ## Check code style and static types
 lint: style static

polygon/rest/aggs.py

@@ -2,6 +2,7 @@
 from typing import Optional, Any, Dict, List, Union
 from .models import Agg, Sort
 from urllib3 import HTTPResponse
+from datetime import datetime
 
 # https://polygon.io/docs/stocks
 class AggsClient(BaseClient):
@@ -11,8 +12,8 @@ def get_aggs(
         multiplier: int,
         timespan: str,
         # "from" is a keyword in python https://www.w3schools.com/python/python_ref_keywords.asp
-        from_: str,
-        to: str,
+        from_: Union[str, int, datetime],
+        to: Union[str, int, datetime],
         adjusted: Optional[bool] = None,
         sort: Optional[Union[str, Sort]] = None,
         limit: Optional[int] = None,
@@ -25,8 +26,8 @@ def get_aggs(
         :param ticker: The ticker symbol.
         :param multiplier: The size of the timespan multiplier.
         :param timespan: The size of the time window.
-        :param _from: The start of the aggregate time window.
-        :param to: The end of the aggregate time window.
+        :param _from: The start of the aggregate time window as YYYY-MM-DD, Unix MS Timestamps, or a datetime.
+        :param to: The end of the aggregate time window as YYYY-MM-DD, Unix MS Timestamps, or a datetime.
         :param adjusted: Whether or not the results are adjusted for splits. By default, results are adjusted. Set this to false to get results that are NOT adjusted for splits.
         :param sort: Sort the results by timestamp. asc will return results in ascending order (oldest at the top), desc will return results in descending order (newest at the top).The end of the aggregate time window.
         :param limit: Limits the number of base aggregates queried to create the aggregate results. Max 50000 and Default 5000. Read more about how limit is used to calculate aggregate results in our article on Aggregate Data API Improvements.
@@ -35,6 +36,11 @@ def get_aggs(
         :return: List of aggregates
         :rtype: List[Agg]
         """
+        if isinstance(from_, datetime):
+            from_ = int(from_.timestamp() * 1000)
+
+        if isinstance(to, datetime):
+            to = int(to.timestamp() * 1000)
         url = f"/v2/aggs/ticker/{ticker}/range/{multiplier}/{timespan}/{from_}/{to}"
 
         return self._get(

polygon/rest/reference.py

@@ -74,7 +74,6 @@ def list_tickers(
     ) -> Union[Iterator[Ticker], HTTPResponse]:
         """
         Query all ticker symbols which are supported by Polygon.io. This API currently includes Stocks/Equities, Crypto, and Forex.
-
         :param ticker: Specify a ticker symbol. Defaults to empty string which queries all tickers.
         :param ticker_lt: Ticker less than
         :param ticker_lte: Ticker less than or equal to
@@ -113,7 +112,6 @@ def get_ticker_details(
     ) -> Union[TickerDetails, HTTPResponse]:
         """
         Get a single ticker supported by Polygon.io. This response will have detailed information about the ticker and the company behind it.
-
         :param ticker: The ticker symbol of the asset.
         :param date: Specify a point in time to get information about the ticker available on that date. When retrieving information from SEC filings, we compare this date with the period of report date on the SEC filing.
         :param params: Any additional query params
@@ -143,7 +141,6 @@ def get_ticker_news(
     ) -> Union[TickerDetails, HTTPResponse]:
         """
         Get the most recent news articles relating to a stock ticker symbol, including a summary of the article and a link to the original source.
-
         :param ticker: Return results that contain this ticker.
         :param published_utc: Return results published on, before, or after this date.
         :param limit: Limit the number of results returned, default is 10 and max is 1000.
@@ -168,7 +165,6 @@ def get_ticker_types(
     ) -> Union[TickerTypes, HTTPResponse]:
         """
         List all ticker types that Polygon.io has.
-
         :param asset_class: Filter by asset class.
         :param locale: Filter by locale.
         :param params: Any additional query params

tests/mocks.py

@@ -1,16 +1,18 @@
 from polygon import RESTClient
 import unittest
-import httpretty
+import httpretty  # type: ignore
 
 mocks = [
     (
         "/v2/aggs/ticker/AAPL/range/1/day/2005-04-01/2005-04-04",
-        '{"ticker":"AAPL","queryCount":2,"resultsCount":2,"adjusted":true,"results":[{"v":6.42646396e+08,"vw":1.469,"o":1.5032,"c":1.4604,"h":1.5064,"l":1.4489,"t":1112331600000,"n":82132},{"v":5.78172308e+08,"vw":1.4589,"o":1.4639,"c":1.4675,"h":1.4754,"l":1.4343,"t":1112587200000,"n":65543}],"status":"OK","request_id":"12afda77aab3b1936c5fb6ef4241ae42","count":2}'
+        '{"ticker":"AAPL","queryCount":2,"resultsCount":2,"adjusted":true,"results":[{"v":6.42646396e+08,"vw":1.469,"o":1.5032,"c":1.4604,"h":1.5064,"l":1.4489,"t":1112331600000,"n":82132},{"v":5.78172308e+08,"vw":1.4589,"o":1.4639,"c":1.4675,"h":1.4754,"l":1.4343,"t":1112587200000,"n":65543}],"status":"OK","request_id":"12afda77aab3b1936c5fb6ef4241ae42","count":2}',
     )
 ]
 
+
 class BaseTest(unittest.TestCase):
     setup = False
+
     def setUp(self):
         if self.setup:
             return
@@ -19,4 +21,3 @@ def setUp(self):
         for m in mocks:
             httpretty.register_uri(httpretty.GET, c.BASE + m[0], m[1])
         self.setup = True
-

tests/test_aggs.py

@@ -2,10 +2,31 @@
 from polygon.rest.models import Agg
 from mocks import BaseTest
 
+
 class AggsTest(BaseTest):
     def test_get_aggs(self):
         c = RESTClient("")
         aggs = c.get_aggs("AAPL", 1, "day", "2005-04-01", "2005-04-04")
-        expected = [Agg(open=1.5032, high=1.5064, low=1.4489, close=1.4604, volume=642646396.0, vwap=1.469, timestamp=1112331600000, transactions=82132), Agg(open=1.4639, high=1.4754, low=1.4343, close=1.4675, volume=578172308.0, vwap=1.4589, timestamp=1112587200000, transactions=65543)]
+        expected = [
+            Agg(
+                open=1.5032,
+                high=1.5064,
+                low=1.4489,
+                close=1.4604,
+                volume=642646396.0,
+                vwap=1.469,
+                timestamp=1112331600000,
+                transactions=82132,
+            ),
+            Agg(
+                open=1.4639,
+                high=1.4754,
+                low=1.4343,
+                close=1.4675,
+                volume=578172308.0,
+                vwap=1.4589,
+                timestamp=1112587200000,
+                transactions=65543,
+            ),
+        ]
         self.assertEqual(aggs, expected)
-