diff --git a/.gitignore b/.gitignore index 2f7f938..7b48391 100644 --- a/.gitignore +++ b/.gitignore @@ -15,4 +15,6 @@ __pycache__ build -dist \ No newline at end of file +dist + +uv.lock \ No newline at end of file diff --git a/dataloader/_env.py b/dataloader/_env.py index 8599858..f96ed94 100644 --- a/dataloader/_env.py +++ b/dataloader/_env.py @@ -15,3 +15,14 @@ def get_dsn() -> str: "Please set it in your .env file or system environment." ) return dsn + + +def get_fred_key() -> str: + """Fetch the FRED Key from environment variables.""" + fred_key = os.getenv("FRED_KEY") + if not fred_key: + raise EnvironmentError( + "Missing FRED_KEY in environment. " + "Please set it in your .env file or system environment." + ) + return fred_key diff --git a/dataloader/_exception.py b/dataloader/_exception.py new file mode 100644 index 0000000..2cc7c48 --- /dev/null +++ b/dataloader/_exception.py @@ -0,0 +1,15 @@ +"""Custom exceptions for dataloader fetch errors.""" + + +class MacroDataFetchError(Exception): + """Raised when a FRED macro data request fails.""" + + def __init__(self, msg: str): + super().__init__(msg) + + +class EquityDataFetchError(Exception): + """Raised when a yfinance equity data request fails.""" + + def __init__(self, msg: str): + super().__init__(msg) diff --git a/dataloader/_fred_yf_client.py b/dataloader/_fred_yf_client.py new file mode 100644 index 0000000..30bf7f7 --- /dev/null +++ b/dataloader/_fred_yf_client.py @@ -0,0 +1,87 @@ +"""Client for fetching macro and equity data via FRED API and yfinance.""" + +import time +import warnings +from typing import ClassVar, Dict, List + +import numpy as np +import pandas as pd +import requests +import yfinance as yf +from attrs import define + +from ._exception import MacroDataFetchError, EquityDataFetchError + +warnings.filterwarnings("ignore") # Ignore yfinance warnings + + +@define +class FredYfClient: + """Client for fetching macro data from FRED and equity data from yfinance.""" + + _fred_key: str + _fred_base: ClassVar[str] = "https://api.stlouisfed.org/fred/series/observations" + + def fred_fn(self, series_id: str, start: str, end: str) -> pd.Series: + """Fetch a single FRED series and return it as a dated pd.Series.""" + params = { + "api_key": self._fred_key, + "series_id": series_id, + "file_type": "json", + "observation_start": start, + "observation_end": end, + } + + try: + req = requests.get(type(self)._fred_base, params=params, timeout=15) + req.raise_for_status() + except requests.exceptions.RequestException as e: + raise MacroDataFetchError(f"HTTP error for FRED series {series_id}") from e + + observations = req.json().get("observations", []) + if not observations: + raise MacroDataFetchError(f"No observations for FRED series {series_id}") + + idx = [pd.to_datetime(obs["date"]) for obs in observations] + vals = [ + float(obs["value"]) if obs["value"] != "." else np.nan + for obs in observations + ] + + return pd.Series(vals, index=idx, name=series_id) + + def fetch_macro_data( + self, macro_tickers: List[str], start: str, end: str + ) -> pd.DataFrame: + """Fetch multiple FRED series and return a combined DataFrame.""" + raw_macro: Dict[str, pd.Series] = {} + + for ticker in macro_tickers: + try: + raw_macro[ticker] = self.fred_fn(ticker, start, end) + except Exception as e: + raise MacroDataFetchError( + f"Failed to fetch macro data for ticker {ticker}" + ) from e + + time.sleep(0.25) + + return pd.DataFrame(raw_macro).sort_index().dropna(how="all") + + def fetch_equity_data( + self, equity_tickers: List[str], start: str, end: str, adj_close: bool = True + ) -> pd.DataFrame: + """Download OHLCV data for equity tickers via yfinance.""" + try: + data = yf.download( + equity_tickers, + start=start, + end=end, + auto_adjust=adj_close, + progress=False, + ) + return data + except Exception as e: + raise EquityDataFetchError( + f"Failed to fetch equity data for tickers: {equity_tickers}" + ) from e diff --git a/dataloader/_manager.py b/dataloader/_manager.py index 15607a4..b3776f2 100644 --- a/dataloader/_manager.py +++ b/dataloader/_manager.py @@ -5,7 +5,8 @@ from attrs import define from ._client import Client -from ._env import get_dsn +from ._fred_yf_client import FredYfClient +from ._env import get_dsn, get_fred_key @define @@ -40,3 +41,7 @@ def close(cls): if cls._instance: cls._instance.close() cls._instance = None + + @classmethod + def get_alt_connection(cls): + return FredYfClient(get_fred_key()) diff --git a/dataloader/loader.py b/dataloader/loader.py index 6bb8367..c9838ba 100644 --- a/dataloader/loader.py +++ b/dataloader/loader.py @@ -1,12 +1,14 @@ """Clickhouse data loader""" -from typing import Any, Dict, List, Optional, ClassVar +from datetime import datetime +from typing import Any, ClassVar, Dict, List, Optional import pandas as pd from attrs import define from ._manager import Manager from ._client import Client +from ._fred_yf_client import FredYfClient @define @@ -14,6 +16,8 @@ class DataLoader: """ A unified data loader for ClickHouse tables and materialized views. + Also supports data fetching via FRED API and Yahoo Finance + Supports: - dynamic column selection - filters and parameters @@ -22,6 +26,34 @@ class DataLoader: database: ClassVar[str] = "ssmif_quant" client: ClassVar[Client] = Manager.get_connection() + alt_client: ClassVar[FredYfClient] = Manager.get_alt_connection() + + @classmethod + def query_fred_yf( + cls, + macro_tickers: List[str], + equity_tickers: List[str], + start_date: str, + end_date: str, + ) -> Dict[str, pd.DataFrame]: + """Fetch macro data from FRED and equity data from yfinance.""" + try: + datetime.strptime(start_date, "%Y-%m-%d") + datetime.strptime(end_date, "%Y-%m-%d") + except ValueError as e: + raise ValueError( + f"Invalid date format. Use YYYY-MM-DD. Got: {start_date}, {end_date}" + ) from e + + dfs = {} + dfs["macro"] = cls.alt_client.fetch_macro_data( + macro_tickers, start_date, end_date + ) + dfs["equity"] = cls.alt_client.fetch_equity_data( + equity_tickers, start_date, end_date + ) + + return dfs @classmethod def query( diff --git a/pyproject.toml b/pyproject.toml index f9d7cca..60500b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,8 @@ dependencies = [ "clickhouse_driver~=0.2.9", "dotenv~=0.9.9", "numpy~=2.3.2", - "pandas~=2.3.2" + "pandas~=2.3.2", + "yfinance>=1.3.0", ] [project.optional-dependencies] @@ -32,6 +33,10 @@ target-version = ["py313"] [tool.pylint] max-line-length = 88 +[[tool.mypy.overrides]] +module = "yfinance" +ignore_missing_imports = true + [tool.setuptools] include-package-data = true diff --git a/requirements.txt b/requirements.txt index b90fad5..7949939 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,8 @@ clickhouse_driver==0.2.9 dotenv==0.9.9 numpy==2.3.2 pandas==2.3.2 +yfinance==1.3.0 setuptools==80.9.0 -wheel==0.45.1 +wheel==0.46.2 twine==6.2.0 build==1.3.0 \ No newline at end of file diff --git a/requirements_dev.txt b/requirements_dev.txt index 77fbe87..267e922 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -1,5 +1,6 @@ -black==25.1.0 +black==26.3.1 mypy==1.17.1 pandas-stubs==2.3.2.250827 pylint==3.3.8 -types-PyYAML==6.0.12.20250915 \ No newline at end of file +types-PyYAML==6.0.12.20250915 +types-requests==2.32.4.20250611 \ No newline at end of file