massive-com · clickingbuttons · Apr 25, 2022 · Apr 25, 2022
diff --git a/polygon/rest/__init__.py b/polygon/rest/__init__.py
@@ -1 +1,6 @@
-from .client import RESTClient
+from .aggs import AggsClient
+from .trades import TradesClient
+
+class RESTClient(AggsClient, TradesClient):
+    pass
+
diff --git a/polygon/rest/aggs.py b/polygon/rest/aggs.py
@@ -0,0 +1,39 @@
+from .base import BaseClient
+from typing import Optional, Any, Dict, List, Union
+from .models import Agg, Sort
+
+# https://polygon.io/docs/stocks
+class AggsClient(BaseClient):
+    def get_aggs(self,
+        ticker: str,
+        multiplier: int,
+        timespan: str,
+        # "from" is a keyword in python https://www.w3schools.com/python/python_ref_keywords.asp
+        from_: str,
+        to: str,
+        adjusted: Optional[bool]=None,
+        sort: Optional[Union[str, Sort]]=None,
+        limit: Optional[int]=None,
+        params: Optional[Dict[str, Any]]=None,
+        raw: bool=False
+    ) -> List[Agg]:
+        """
+        Get aggregate bars for a ticker over a given date range in custom time window sizes.
+
+        :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 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.
+        :param params: Any additional query params
+        :param raw: Return raw object instead of results object
+        :return: List of aggregates
+        :rtype: List[Agg]
+        """
+        url = f"/v2/aggs/ticker/{ticker}/range/{multiplier}/{timespan}/{from_}/{to}"
+
+        return self._get(path=url, params=self._get_params(self.get_aggs, locals()), resultKey="results", deserializer=Agg.from_dict, raw=raw)
+
diff --git a/polygon/rest/base.py b/polygon/rest/base.py
@@ -0,0 +1,90 @@
+import os
+import json
+import urllib3
+import inspect
+from enum import Enum
+from typing import Optional, Any
+
+base = 'https://api.polygon.io'
+env_key = "POLYGON_API_KEY"
+
+# https://urllib3.readthedocs.io/en/stable/reference/urllib3.poolmanager.html
+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
+        ):
+        if api_key is None:
+            raise Exception(f"Must specify env var {env_key} or pass api_key in constructor")
+        self.API_KEY = api_key
+        self.BASE = base
+
+        # https://urllib3.readthedocs.io/en/stable/reference/urllib3.connectionpool.html#urllib3.HTTPConnectionPool
+        self.client = urllib3.PoolManager(num_pools=num_pools, headers={
+            'Authorization': 'Bearer ' + self.API_KEY
+        })
+        self.timeout=urllib3.Timeout(connect=connect_timeout, read=read_timeout)
+        self.retries = retries
+
+    def _decode(self, resp):
+        return json.loads(resp.data.decode('utf-8'))
+
+    def _get(self, path: str, params: Optional[dict] = None, resultKey: Optional[str] = None, deserializer = None, raw: bool = False) -> Any:
+        if params is None:
+            params = {}
+        params = {str(k): str(v) for k, v in params.items() if v is not None}
+        resp = self.client.request('GET', self.BASE + path, fields=params, retries=self.retries)
+
+        if resp.status != 200:
+            raise Exception(resp.data.decode('utf-8'))
+
+        if raw:
+            return resp
+
+        obj = self._decode(resp)
+
+        if resultKey:
+            obj = obj[resultKey]
+
+        if deserializer:
+            obj = [deserializer(o) for o in obj]
+
+        return obj
+
+    def _get_params(self, fn, caller_locals):
+        params = caller_locals["params"]
+        if params is None:
+            params = {}
+        # https://docs.python.org/3.7/library/inspect.html#inspect.Signature
+        for argname, v in inspect.signature(fn).parameters.items():
+            # https://docs.python.org/3.7/library/inspect.html#inspect.Parameter
+            if argname in ['params', 'raw']:
+                continue
+            if v.default != v.empty:
+                # timestamp_lt -> timestamp.lt
+                val = caller_locals.get(argname, v.default)
+                if isinstance(val, Enum):
+                    val = val.value
+                if val is not None:
+                    params[argname.replace("_", ".")] = val
+
+        return params
+
+    def _paginate(self, path: str, params: dict, raw: bool, deserializer):
+        while True:
+            resp = self._get(path=path, params=params, deserializer=deserializer, raw=True)
+            if raw:
+                return resp
+            decoded = self._decode(resp)
+            for t in decoded["results"]:
+                yield deserializer(t)
+            if "next_url" in decoded:
+                path = decoded["next_url"].replace(self.BASE, '')
+                params = {}
+            else:
+                return