"""Data pipeline — fetch OHLCV data from yfinance, cache to disk."""
import json
import os
import pickle
import logging
from datetime import datetime
from typing import Dict, List, Optional, Tuple
import pandas as pd
import yfinance as yf
from backend.config import config
logger = logging.getLogger(__name__)
PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
def _data_path() -> str:
path = os.path.join(PROJECT_ROOT, config.storage.get("data_dir", "data/cache"))
os.makedirs(path, exist_ok=True)
return path
def _ticker_file(ticker: str) -> str:
return os.path.join(_data_path(), f"{ticker}_raw.parquet")
def _cache_file(ticker: str) -> str:
return os.path.join(_data_path(), f"{ticker}_meta.json")
def fetch_ticker(ticker: str, interval: str = None, period: str = "max") -> pd.DataFrame:
"""Fetch OHLCV data for a ticker using yfinance.
Returns DataFrame with index as datetime, columns: Open, High, Low, Close, Volume.
Note: yfinance limits 4h/1h data to 730 days. For longer history, use 1d interval.
"""
# Load interval from config if not specified
if interval is None:
from backend.config import config
interval = config.data.get("interval", "1d")
df = yf.download(ticker, interval=interval, period=period, progress=False)
# Flatten multi-level columns (yfinance quirk)
if isinstance(df.columns, pd.MultiIndex):
df.columns = df.columns.get_level_values(0)
# Clean up
df = df.dropna(how="all")
if df.empty:
logger.warning(f"No data returned for {ticker} (interval={interval}, period={period})")
return df
df.index.name = "Date"
df.index = pd.to_datetime(df.index)
# Sort by date
df = df.sort_index()
logger.info(f"Fetched {len(df)} bars for {ticker} ({df.index[0]} → {df.index[-1]})")
return df
def fetch_all_tickers() -> Dict[str, pd.DataFrame]:
"""Fetch all configured tickers. Returns dict of ticker -> DataFrame."""
tickers = config.data.get("tickers", ["AAPL"])
results = {}
errors = []
for ticker in tickers:
try:
df = fetch_ticker(ticker)
if not df.empty:
results[ticker] = df
# Cache
_cache_ticker(ticker, df)
logger.info(f" ✓ {ticker}: {len(df)} bars")
else:
errors.append(f"{ticker}: empty data")
except Exception as e:
errors.append(f"{ticker}: {e}")
logger.error(f" ✗ {ticker}: {e}")
if errors:
logger.warning(f"Errors: {', '.join(errors)}")
return results
def _cache_ticker(ticker: str, df: pd.DataFrame):
"""Save raw data to parquet + metadata."""
path = _ticker_file(ticker)
df.to_parquet(path, index=True)
meta = {
"ticker": ticker,
"rows": len(df),
"start": str(df.index[0]),
"end": str(df.index[-1]),
"fetched_at": datetime.now().isoformat(),
"columns": list(df.columns),
}
with open(_cache_file(ticker), "w") as f:
json.dump(meta, f, indent=2)
def load_cached(ticker: str) -> Optional[pd.DataFrame]:
"""Load cached parquet data for a ticker."""
path = _ticker_file(ticker)
if os.path.exists(path):
return pd.read_parquet(path)
return None
def load_all_cached() -> Dict[str, pd.DataFrame]:
"""Load all cached ticker data."""
results = {}
data_dir = _data_path()
for fname in os.listdir(data_dir):
if fname.endswith("_raw.parquet"):
ticker = fname.replace("_raw.parquet", "")
try:
results[ticker] = pd.read_parquet(os.path.join(data_dir, fname))
except Exception as e:
logger.warning(f"Failed to load {ticker}: {e}")
return results
def get_latest_bar(ticker: str) -> Optional[dict]:
"""Get the most recent OHLCV bar for a ticker."""
df = load_cached(ticker)
if df is None or df.empty:
df = fetch_ticker(ticker)
if df.empty:
return None
_cache_ticker(ticker, df)
last = df.iloc[-1]
return {
"date": str(df.index[-1]),
"open": float(last.get("Open", 0)),
"high": float(last.get("High", 0)),
"low": float(last.get("Low", 0)),
"close": float(last.get("Close", 0)),
"volume": int(last.get("Volume", 0)),
}