add better types (#118)

Author: clickingbuttons

polygon/rest/aggs.py

@@ -1,6 +1,7 @@
 from .base import BaseClient
 from typing import Optional, Any, Dict, List, Union
 from .models import Agg, Sort
+from urllib3 import HTTPResponse
 
 # https://polygon.io/docs/stocks
 class AggsClient(BaseClient):
@@ -17,7 +18,7 @@ def get_aggs(
         limit: Optional[int] = None,
         params: Optional[Dict[str, Any]] = None,
         raw: bool = False,
-    ) -> List[Agg]:
+    ) -> Union[List[Agg], HTTPResponse]:
         """
         Get aggregate bars for a ticker over a given date range in custom time window sizes.
 

polygon/rest/models/aggs.py

@@ -1,17 +1,26 @@
 from dataclasses import dataclass
-
+from typing import Optional
 
 @dataclass
 class Agg:
+    timestamp: int
     open: float
     high: float
     low: float
     close: float
     volume: float
-    vwap: float
-    timestamp: int
-    transactions: int
+    vwap: Optional[float]
+    transactions: Optional[int]
 
     @staticmethod
     def from_dict(d):
-        return Agg(d["o"], d["h"], d["l"], d["c"], d["v"], d["vw"], d["t"], d["n"])
+        return Agg(
+                timestamp=d["t"],
+                open=d["o"],
+                high=d["h"], 
+                low=d["l"],
+                close=d["c"],
+                volume=d["v"],
+                vwap=d.get("vw", None),
+                transactions=d.get("n", None)
+        )

polygon/rest/trades.py

@@ -1,6 +1,7 @@
 from .base import BaseClient
-from typing import Optional, Any, Dict, Union
+from typing import Optional, Any, Dict, Union, Iterator
 from .models import Trade, Sort, Order
+from urllib3 import HTTPResponse
 
 # https://polygon.io/docs/stocks
 class TradesClient(BaseClient):
@@ -17,7 +18,7 @@ def list_trades(
         order: Optional[Union[str, Order]] = None,
         params: Optional[Dict[str, Any]] = None,
         raw: bool = False,
-    ):
+    ) -> Union[Iterator[Trade], HTTPResponse]:
         """
         Get trades for a ticker symbol in a given time range.