mbleh

paperone.py

@@ -8,15 +8,18 @@ from dotenv import load_dotenv
 from typing import NoReturn
 from os import environ
 from enum import Enum
+from dataclasses import dataclass
 
 load_dotenv()
 
-API_KEY: str | None = environ.get("API_KEY")
-BASE_URL: str = "https://api.taapi.io"
+
+@dataclass(frozen=True)
+class Indicator:
+    endpoint: str
 
 
-class Indicator(Enum):
-    rsi = "rsi"
+class IndicatorEnum(Enum):
+    RSI = Indicator("rsi")
 
 
 class Interval(Enum):
@@ -32,42 +35,43 @@ class Interval(Enum):
     OneWeek = 604800
 
 
+class TaapiClient:
+    def __init__(self, api_key: str) -> None:
+        self._api_key: str = api_key
+        self._base_url: str = "https://api.taapi.io"
+        self._session: requests.Session = requests.Session()
+
+    def __build_indicator_url(self, indicator: Indicator) -> str:
+        return f"{self._base_url}/{indicator.endpoint}"
+
     def query_indicator(
-    ticker: str, indicator: Indicator, target_date: datetime
+        self,
+        ticker: str,
+        indicator: Indicator,
+        target_date: datetime,
+        interval: str = "1d",
+        results: int = 14,
     ) -> requests.Response:
-    interval: str = "1d"
-    candles: int = calculate_candles_to_target_date(target_date, interval)
-    target_url = f"{BASE_URL}/{indicator}"
+        backtrack_candles: int = self._candles_to_target_date(target_date, interval)
+        target_url: str = self.__build_indicator_url(indicator)
 
         params: dict[str, str | int | bool] = {
-        "secret": API_KEY,
+            "secret": self._api_key,
             "symbol": ticker,
             "interval": interval,
             "type": "stocks",
             "gaps": "false",
             "addResultTimestamp": "true",
-        "backtrack": candles,
-        "results": "14",
+            "backtrack": backtrack_candles,
+            "results": str(results),
         }
 
-    response: requests.Response = requests.get(target_url, params=params)
+        response: requests.Response = self._session.get(target_url, params=params)
 
         return response
 
-
-def is_trading_day(date: datetime) -> bool:
-    return date.weekday() not in [5, 6]
-
-
-def parse_date_yyyymmdd(date_str: str) -> datetime:
-    return datetime.strptime(date_str, "%Y%m%d")
-
-
-def format_date_readable(date: datetime) -> str:
-    return date.strftime("%B %d, %Y")
-
-
-def calculate_candles_to_target_date(
+    @staticmethod
+    def _candles_to_target_date(
         target_date: datetime,
         interval: str = "1h",
         current_time: datetime | None = None,
@@ -100,17 +104,47 @@ def calculate_candles_to_target_date(
 
         return num_candles
 
+    def close(self) -> None:
+        self._session.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
+
+def is_trading_day(date: datetime) -> bool:
+    return date.weekday() not in [5, 6]
+
+
+def parse_date_yyyymmdd(date_str: str) -> datetime:
+    return datetime.strptime(date_str, "%Y%m%d")
+
+
+def format_date_readable(date: datetime) -> str:
+    return date.strftime("%B %d, %Y")
+
 
 def main() -> NoReturn:
+    api_key = environ.get("API_KEY")
+
+    if not api_key:
+        print("API_KEY not set")
+
+        exit(0)
+
     date = parse_date_yyyymmdd("20250821")
-    indicator = Indicator.rsi
+    indicator: Indicator = IndicatorEnum.RSI.value
     ticker = "AAPL"
-    response = query_indicator(ticker, indicator.value, date)
+
+    with TaapiClient(api_key) as client:
+        response = client.query_indicator(ticker, indicator, date)
 
         if response.status_code == 200:
             data: dict[str, list[float] | list[int]] = response.json()
 
-        print(f"{indicator.value} values for {ticker} on {format_date_readable(date)}:")
+            print(f"{indicator} values for {ticker} on {format_date_readable(date)}:")
 
             for rsi_val, ts in zip(data["value"], data["timestamp"]):
                 dt: datetime = datetime.fromtimestamp(ts)