Warrior_EA/research/altdata/common.py

"""Shared plumbing for alternative-data collectors.

Every collector writes rows that carry TWO timestamps:
  observed   - the period the value describes (e.g. COT report Tuesday)
  published  - when the value became publicly knowable (UTC)
Joins against price bars must use `published <= bar_open`, never `observed`.
This is the vintage discipline; see DESIGN.md.
"""
import io
import json
import os
import time
import zipfile
from pathlib import Path

import requests

DATA_ROOT = Path(r"c:\Users\admin\Documents\Workspaces\Market Data\altdata")
RAW = DATA_ROOT / "raw"
KEYS_FILE = DATA_ROOT / "keys.json"

_HEADERS = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) Warrior_EA-research/1.0"}


def ensure_dirs() -> None:
    RAW.mkdir(parents=True, exist_ok=True)


def api_key(name: str) -> str | None:
    """Key lookup: env ALTDATA_<NAME>_KEY first, then Market Data/altdata/keys.json."""
    env = os.environ.get(f"ALTDATA_{name.upper()}_KEY")
    if env:
        return env
    if KEYS_FILE.exists():
        try:
            return json.loads(KEYS_FILE.read_text()).get(name.lower())
        except (json.JSONDecodeError, OSError):
            pass
    return None


def get(url: str, *, params: dict | None = None, retries: int = 3, timeout: int = 60,
        ok_codes: tuple[int, ...] = (200,)) -> requests.Response:
    last = None
    for attempt in range(retries):
        try:
            r = requests.get(url, params=params, headers=_HEADERS, timeout=timeout)
            if r.status_code in ok_codes:
                return r
            last = RuntimeError(f"HTTP {r.status_code} for {url}")
        except requests.RequestException as e:  # noqa: PERF203
            last = e
        time.sleep(2 * (attempt + 1))
    raise RuntimeError(f"GET failed after {retries} tries: {url}") from last


def cached_download(url: str, dest: Path, *, refresh: bool = False) -> Path:
    """Download url to dest unless it already exists. Atomic write (tmp then rename)."""
    if dest.exists() and not refresh:
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    r = get(url)
    tmp = dest.with_suffix(dest.suffix + ".tmp")
    tmp.write_bytes(r.content)
    tmp.replace(dest)
    return dest


def read_zip_member(zip_path: Path) -> tuple[str, bytes]:
    """Return (name, bytes) of the single data member of a CFTC-style zip."""
    with zipfile.ZipFile(io.BytesIO(zip_path.read_bytes())) as z:
        names = [n for n in z.namelist() if not n.endswith("/")]
        if len(names) != 1:
            raise ValueError(f"{zip_path}: expected 1 member, got {names}")
        return names[0], z.read(names[0])