src/option_implied_lab/data.py

from __future__ import annotations

from datetime import datetime, timezone

import pandas as pd
import yfinance as yf

from .types import OptionSnapshot


def _normalize_chain(
    df: pd.DataFrame, option_type: str, expiry: pd.Timestamp, spot: float
) -> pd.DataFrame:
    out = df.copy()
    out["option_type"] = option_type
    out["expiry"] = pd.to_datetime(expiry).tz_localize(None)
    out["mid"] = (out["bid"].fillna(0.0) + out["ask"].fillna(0.0)) / 2.0
    out["volume"] = out.get("volume", 0.0)
    out["openInterest"] = out.get("openInterest", 0.0)
    out = out[
        [
            "expiry",
            "option_type",
            "strike",
            "bid",
            "ask",
            "mid",
            "volume",
            "openInterest",
        ]
    ].copy()
    out = out.dropna(subset=["strike", "mid"])
    out = out[out["strike"] > 0].copy()
    out["moneyness"] = out["strike"] / float(spot)
    return out


def fetch_option_snapshot(ticker: str, max_expiries: int = 2) -> OptionSnapshot:
    tk = yf.Ticker(ticker)
    hist = tk.history(period="1d")
    if hist.empty:
        raise ValueError(f"No price history for ticker {ticker}")
    spot = float(hist["Close"].iloc[-1])

    expiries = tk.options[:max_expiries]
    if not expiries:
        raise ValueError(f"No option expiries for ticker {ticker}")

    rows = []
    for expiry_str in expiries:
        chain = tk.option_chain(expiry_str)
        expiry = pd.to_datetime(expiry_str)
        rows.append(_normalize_chain(chain.calls, "call", expiry, spot))
        rows.append(_normalize_chain(chain.puts, "put", expiry, spot))

    options = pd.concat(rows, ignore_index=True)
    options = options.sort_values(["expiry", "option_type", "strike"]).reset_index(
        drop=True
    )

    return OptionSnapshot(
        ticker=ticker,
        snapshot_time=datetime.now(timezone.utc),
        spot=spot,
        options=options,
    )