Warrior_EA/research/altdata/eia.py

"""EIA v2 API collector - free key (https://www.eia.gov/opendata/register.php).

Weekly Petroleum Status Report series for XTIUSD:
    crude stocks (ex-SPR), field production, refinery utilization, SPR stocks.

PUBLICATION LAG: the WPSR covers the week ending FRIDAY and is released the
following WEDNESDAY 10:30 ET (Thursday if a federal holiday intervenes).
We stamp `published` = Thursday 00:00 UTC after the covering week - one day
conservative, holiday-safe, and correct for D1 usage.

Usage:
    python -m altdata.eia
Key: env ALTDATA_EIA_KEY or "eia" in Market Data/altdata/keys.json.
"""
import pandas as pd

from .common import DATA_ROOT, api_key, ensure_dirs, get

API = "https://api.eia.gov/v2/petroleum/sum/sndw/data/"

SERIES = {
    "WCESTUS1": "crude_stocks_ex_spr",
    "WCRFPUS2": "crude_field_production",
    "WPULEUS3": "refinery_utilization_pct",
    "WCSSTUS1": "spr_stocks",
}


def fetch(key: str) -> pd.DataFrame:
    rows, offset = [], 0
    while True:
        params = {
            "api_key": key, "frequency": "weekly",
            "data[0]": "value", "offset": offset, "length": 5000,
        }
        for i, sid in enumerate(SERIES):
            params[f"facets[series][{i}]"] = sid
        j = get(API, params=params).json()["response"]
        rows.extend(j["data"])
        offset += 5000
        if offset >= int(j["total"]):
            break
    df = pd.DataFrame(rows)
    df["observed"] = pd.to_datetime(df["period"])
    df["value"] = pd.to_numeric(df["value"], errors="coerce")
    # Wednesday after the covering week ends Friday; +1 day guard = Thursday 00:00 UTC
    df["published"] = df["observed"] + pd.to_timedelta(6, unit="D")
    return df


def main() -> None:
    key = api_key("eia")
    if not key:
        raise SystemExit("No EIA key. Set ALTDATA_EIA_KEY or add 'eia' to "
                         f"{DATA_ROOT / 'keys.json'} (free: eia.gov/opendata)")
    ensure_dirs()
    out_dir = DATA_ROOT / "eia"
    out_dir.mkdir(parents=True, exist_ok=True)
    df = fetch(key)
    for sid, name in SERIES.items():
        sub = df[df["series"] == sid][["observed", "published", "value"]]
        sub = sub.sort_values("observed").reset_index(drop=True)
        dest = out_dir / f"{name}.csv"
        sub.to_csv(dest, index=False)
        print(f"{sid} ({name}): {len(sub)} weekly rows => {dest.name}")


if __name__ == "__main__":
    main()