aggs endpoints (#128)

mcdayoub · clickingbuttons · commit 56b4a04699ea · 2022-04-29T15:47:21.000-04:00
* aggs endpoints

* aggs endpoints

* aggs models optional

* aggs models optional

* fixed result_keys
diff --git a/polygon/rest/aggs.py b/polygon/rest/aggs.py
@@ -1,6 +1,7 @@
+from email.headerregistry import Group
 from .base import BaseClient
 from typing import Optional, Any, Dict, List, Union
-from .models import Agg, Sort
+from .models import Agg, GroupedDailyAgg, DailyOpenCloseAgg, PreviousCloseAgg, Sort
 from urllib3 import HTTPResponse
 from datetime import datetime
 
@@ -22,7 +23,6 @@ def get_aggs(
     ) -> Union[List[Agg], HTTPResponse]:
         """
         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.
@@ -34,7 +34,6 @@ def get_aggs(
         :param params: Any additional query params
         :param raw: Return raw object instead of results object
         :return: List of aggregates
-        :rtype: List[Agg]
         """
         if isinstance(from_, datetime):
             from_ = int(from_.timestamp() * 1000)
@@ -50,3 +49,82 @@ def get_aggs(
             deserializer=Agg.from_dict,
             raw=raw,
         )
+
+    def get_grouped_daily_aggs(
+        self,
+        date: str,
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+    ) -> Union[List[GroupedDailyAgg], HTTPResponse]:
+        """
+        Get the daily open, high, low, and close (OHLC) for the entire market.
+
+        :param date: The beginning date for the aggregate 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 params: Any additional query params
+        :param raw: Return raw object instead of results object
+        :return: List of grouped daily aggregates
+        """
+        url = f"/v2/aggs/grouped/locale/us/market/stocks/{date}"
+
+        return self._get(
+            path=url,
+            params=self._get_params(self.get_grouped_daily_aggs, locals()),
+            result_key="results",
+            deserializer=GroupedDailyAgg.from_dict,
+            raw=raw,
+        )
+
+    def get_daily_open_close_agg(
+        self,
+        ticker: str,
+        date: str,
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+    ) -> Union[DailyOpenCloseAgg, HTTPResponse]:
+        """
+        Get the open, close and afterhours prices of a stock symbol on a certain date.
+
+        :param ticker: The exchange symbol that this item is traded under.
+        :param date: The beginning date for the aggregate 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 params: Any additional query params
+        :param raw: Return raw object instead of results object
+        :return: Daily open close aggregate
+        """
+        url = f"/v1/open-close/{ticker}/{date}"
+
+        return self._get(
+            path=url,
+            params=self._get_params(self.get_daily_open_close_agg, locals()),
+            deserializer=DailyOpenCloseAgg.from_dict,
+            raw=raw,
+        )
+
+    def get_previous_close_agg(
+        self,
+        ticker: str,
+        adjusted: Optional[bool] = None,
+        params: Optional[Dict[str, Any]] = None,
+        raw: bool = False,
+    ) -> Union[PreviousCloseAgg, HTTPResponse]:
+        """
+        Get the previous day's open, high, low, and close (OHLC) for the specified stock ticker.
+
+        :param ticker: The ticker symbol of the stock/equity.
+        :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 params: Any additional query params
+        :param raw: Return raw object instead of results object
+        :return: Previous close aggregate
+        """
+        url = f"/v2/aggs/ticker/{ticker}/prev"
+
+        return self._get(
+            path=url,
+            params=self._get_params(self.get_previous_close_agg, locals()),
+            result_key="results",
+            deserializer=PreviousCloseAgg.from_dict,
+            raw=raw,
+        )
diff --git a/polygon/rest/base.py b/polygon/rest/base.py
@@ -61,6 +61,9 @@ def _get(
 
         if result_key:
             obj = obj[result_key]
+        else:
+            # If the result_key does not exist, still need to put the results in a list
+            obj = [obj]
 
         if deserializer:
             obj = [deserializer(o) for o in obj]
diff --git a/polygon/rest/models/aggs.py b/polygon/rest/models/aggs.py
@@ -4,24 +4,105 @@
 
 @dataclass
 class Agg:
-    timestamp: int
     open: float
     high: float
     low: float
     close: float
     volume: float
     vwap: Optional[float]
+    timestamp: Optional[int]
     transactions: Optional[int]
 
     @staticmethod
     def from_dict(d):
         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),
+            d.get("o", None),
+            d.get("h", None),
+            d.get("l", None),
+            d.get("c", None),
+            d.get("v", None),
+            d.get("vw", None),
+            d.get("t", None),
+            d.get("n", None),
+        )
+
+
+@dataclass
+class GroupedDailyAgg:
+    ticker: str
+    open: float
+    high: float
+    low: float
+    close: float
+    volume: float
+    vwap: Optional[float]
+    timestamp: Optional[int]
+    transactions: Optional[int]
+
+    @staticmethod
+    def from_dict(d):
+        return GroupedDailyAgg(
+            d.get("T", None),
+            d.get("o", None),
+            d.get("h", None),
+            d.get("l", None),
+            d.get("c", None),
+            d.get("v", None),
+            d.get("vw", None),
+            d.get("t", None),
+            d.get("n", None),
+        )
+
+
+@dataclass
+class DailyOpenCloseAgg:
+    after_hours: Optional[float]
+    close: float
+    from_: str
+    high: float
+    low: float
+    open: float
+    pre_market: Optional[float]
+    status: Optional[str]
+    symbol: str
+    volume: float
+
+    @staticmethod
+    def from_dict(d):
+        return DailyOpenCloseAgg(
+            d.get("afterHours", None),
+            d.get("close", None),
+            d.get("from", None),
+            d.get("high", None),
+            d.get("low", None),
+            d.get("open", None),
+            d.get("preMarket", None),
+            d.get("status", None),
+            d.get("symbol", None),
+            d.get("volume", None),
+        )
+
+
+@dataclass
+class PreviousCloseAgg:
+    ticker: str
+    close: float
+    high: float
+    low: float
+    open: float
+    timestamp: Optional[float]
+    volume: float
+    vwap: Optional[float]
+
+    @staticmethod
+    def from_dict(d):
+        return PreviousCloseAgg(
+            d.get("T", None),
+            d.get("c", None),
+            d.get("h", None),
+            d.get("l", None),
+            d.get("o", None),
+            d.get("t", None),
+            d.get("v", None),
+            d.get("vw", None),
         )
diff --git a/tests/mocks.py b/tests/mocks.py
@@ -6,7 +6,19 @@
     (
         "/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}',
-    )
+    ),
+    (
+        "/v2/aggs/grouped/locale/us/market/stocks/2005-04-04",
+        '{"queryCount":1,"resultsCount":1,"adjusted": true,"results": [{"T":"GIK","v":895345,"vw":9.9979,"o":9.99,"c":10.02,"h":10.02,"l":9.9,"t":1602705600000,"n":96}],"status":"OK","request_id":"eae3ded2d6d43f978125b7a8a609fad9","count":1}',
+    ),
+    (
+        "/v1/open-close/AAPL/2005-04-01",
+        '{"status": "OK","from": "2021-04-01","symbol": "AAPL","open": 123.66,"high": 124.18,"low": 122.49,"close": 123,"volume": 75089134,"afterHours": 123,"preMarket": 123.45}',
+    ),
+    (
+        "/v2/aggs/ticker/AAPL/prev",
+        '{"ticker":"AAPL","queryCount":1,"resultsCount":1,"adjusted":true,"results":[{"T":"AAPL","v":9.5595226e+07,"vw":158.6074,"o":162.25,"c":156.8,"h":162.34,"l":156.72,"t":1651003200000,"n":899965}],"status":"OK","request_id":"5e5378d5ecaf3df794bb52e45d015d2e","count":1}',
+    ),
 ]
 
 
diff --git a/tests/test_aggs.py b/tests/test_aggs.py
@@ -1,5 +1,10 @@
 from polygon import RESTClient
-from polygon.rest.models import Agg
+from polygon.rest.models import (
+    Agg,
+    GroupedDailyAgg,
+    DailyOpenCloseAgg,
+    PreviousCloseAgg,
+)
 from mocks import BaseTest
 
 
@@ -30,3 +35,57 @@ def test_get_aggs(self):
             ),
         ]
         self.assertEqual(aggs, expected)
+
+    def test_get_grouped_daily_aggs(self):
+        c = RESTClient("")
+        aggs = c.get_grouped_daily_aggs("2005-04-04", True)
+        expected = [
+            GroupedDailyAgg(
+                ticker="GIK",
+                open=9.99,
+                high=10.02,
+                low=9.9,
+                close=10.02,
+                volume=895345,
+                vwap=9.9979,
+                timestamp=1602705600000,
+                transactions=96,
+            )
+        ]
+        self.assertEqual(aggs, expected)
+
+    def test_get_daily_open_close_agg(self):
+        c = RESTClient("")
+        aggs = c.get_daily_open_close_agg("AAPL", "2005-04-01", True)
+        expected = [
+            DailyOpenCloseAgg(
+                after_hours=123,
+                close=123,
+                from_="2021-04-01",
+                high=124.18,
+                low=122.49,
+                open=123.66,
+                pre_market=123.45,
+                status="OK",
+                symbol="AAPL",
+                volume=75089134,
+            )
+        ]
+        self.assertEqual(aggs, expected)
+
+    def test_get_previous_close_agg(self):
+        c = RESTClient("")
+        aggs = c.get_previous_close_agg("AAPL")
+        expected = [
+            PreviousCloseAgg(
+                ticker="AAPL",
+                close=156.8,
+                high=162.34,
+                low=156.72,
+                open=162.25,
+                timestamp=1651003200000,
+                volume=95595226.0,
+                vwap=158.6074,
+            )
+        ]
+        self.assertEqual(aggs, expected)
