"""yfinance API helper functions."""
import logging
import sys
from typing import Literal
import pandas as pd
import yfinance as yf
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception,
after_log,
)
from yfinance.exceptions import YFRateLimitError
# Setup logging
logger = logging.getLogger(__name__)
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
handlers=[logging.StreamHandler(sys.stderr)],
)
def api_retry(func):
"""Unified retry decorator for API calls (yfinance and HTTP).
Retries on:
- YFRateLimitError
- HTTP 5xx errors
- Network/connection issues
- Rate limiting (429)
"""
return retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=2.0, min=2.0, max=30.0),
retry=retry_if_exception(
lambda e: isinstance(e, YFRateLimitError)
or (hasattr(e, "status_code") and getattr(e, "status_code", 0) >= 500)
or any(
term in str(e).lower()
for term in [
"rate limit",
"too many requests",
"temporarily blocked",
"timeout",
"connection",
"network",
"temporary",
"5",
"429",
"502",
"503",
"504",
]
)
),
after=after_log(logger, logging.WARNING),
)(func)
@api_retry
def yf_call(ticker: str, method: str, *args, **kwargs):
"""Generic yfinance API call with retry logic.
Args:
ticker: Stock ticker symbol
method: Method name to call on yf.Ticker object
*args: Positional arguments to pass to the method
**kwargs: Keyword arguments to pass to the method
Returns:
Result of the yfinance method call
"""
t = yf.Ticker(ticker)
return getattr(t, method)(*args, **kwargs)
def get_options_chain(
ticker: str, expiry: str, option_type: Literal["C", "P"] | None = None
) -> pd.DataFrame:
"""Get options chain with optional filtering by type.
Args:
ticker: Stock ticker symbol
expiry: Expiration date
option_type: Option type - "C" for calls, "P" for puts, None for both
Returns:
DataFrame with options chain data
"""
chain = yf_call(ticker, "option_chain", expiry)
if option_type == "C":
return chain.calls
elif option_type == "P":
return chain.puts
return pd.concat([chain.calls, chain.puts], ignore_index=True)
