bls-data/bls_client/client.py

"""
BLSClient — thin wrapper around the BLS Public Data API v2.

Usage:
    from bls_client import BLSClient
    client = BLSClient("YOUR_API_KEY")
    data = client.fetch(["CES0000000001"], 2020, 2025)
"""

import requests
from datetime import datetime

from .cache import FileCache
from .errors import BLSError, BLSRequestError, BLSQuotaError
from .retry import with_retries


class BLSClient:
    BASE_URL = "https://api.bls.gov/publicAPI/v2"
    MAX_SERIES_PER_CALL = 50
    MAX_YEARS_PER_CALL  = 20

    def __init__(self, api_key: str, cache=False, cache_dir=None, cache_ttl=86400,
                 retries=2, backoff=0.5):
        """
        Args:
            api_key:   BLS registration key.
            cache:     True to enable the default on-disk cache (~/.cache/bls),
                       or pass a custom cache object with get(key)/set(key, value).
            cache_dir: Override the default cache directory.
            cache_ttl: Cache entry lifetime in seconds (default 1 day).
            retries:   retries after the first attempt on transient network errors
                       (timeouts, connection drops, 429/5xx). 0 disables retrying.
            backoff:   base exponential-backoff interval in seconds.
        """
        self.api_key = api_key
        if cache is True:
            self.cache = FileCache(cache_dir, cache_ttl)
        elif cache:
            self.cache = cache            # caller-supplied cache object
        else:
            self.cache = None
        self.retries = retries
        self.backoff = backoff
        self.queries_used = 0             # network POSTs spent this session (cache hits don't count)

    def _get(self, url, **kwargs):
        """GET with retry/backoff."""
        return with_retries(lambda: requests.get(url, **kwargs),
                            retries=self.retries, backoff=self.backoff)

    # ------------------------------------------------------------------
    # Core fetch
    # ------------------------------------------------------------------
    def fetch(
        self,
        series_ids: list[str] | str,
        start_year: int,
        end_year:   int,
        catalog:    bool = True,
        calculations: bool = False,
        annual_average: bool = False,
    ) -> dict[str, list[dict]]:
        """
        Fetch time-series data for one or more series IDs.

        Automatically batches requests when series_ids > 50.
        Returns a dict keyed by series ID, value is the list of
        observations (newest first).

        Args:
            series_ids:    Single series ID or list of series IDs.
            start_year:    First year to retrieve.
            end_year:      Last year to retrieve (max 20 years from start).
            catalog:       Include series title/metadata in response.
            calculations:  Include MoM / YoY net and pct changes.
            annual_average: Include M13 annual average rows.

        Returns:
            {series_id: [{"year":..., "period":..., "value":..., ...}, ...]}
        """
        if isinstance(series_ids, str):
            series_ids = [series_ids]

        results = {}
        for i in range(0, len(series_ids), self.MAX_SERIES_PER_CALL):
            batch = series_ids[i : i + self.MAX_SERIES_PER_CALL]
            payload = {
                "seriesid":       batch,
                "startyear":      str(start_year),
                "endyear":        str(end_year),
                "registrationkey": self.api_key,
                "catalog":        catalog,
                "calculations":   calculations,
                "annualaverage":  annual_average,
            }

            cache_key = self.cache.key(payload) if self.cache else None
            batch_result = self.cache.get(cache_key) if cache_key else None

            if batch_result is None:
                batch_result = self._fetch_batch(payload)
                if cache_key:
                    self.cache.set(cache_key, batch_result)

            results.update(batch_result)

        return results

    def _fetch_batch(self, payload: dict) -> dict:
        """Execute one (uncached) API call and return {series_id: {data, catalog}}."""
        url = f"{self.BASE_URL}/timeseries/data/"
        r = with_retries(lambda: requests.post(url, json=payload, timeout=30),
                         retries=self.retries, backoff=self.backoff)
        r.raise_for_status()
        self.queries_used += 1
        body = r.json()

        status = body.get("status")
        if status != "REQUEST_SUCCEEDED":
            msgs = body.get("message", [])
            text = " ".join(msgs) if isinstance(msgs, list) else str(msgs)
            low = text.lower()
            if status == "REQUEST_NOT_PROCESSED" and ("threshold" in low or "exceed" in low):
                raise BLSQuotaError(
                    f"BLS daily query limit reached: {text}", status=status, messages=msgs
                )
            raise BLSRequestError(
                f"BLS API request failed ({status}): {text}", status=status, messages=msgs
            )

        return {
            s["seriesID"]: {"data": s["data"], "catalog": s.get("catalog", {})}
            for s in body["Results"]["series"]
        }

    def fetch_latest(
        self,
        series_ids: list[str] | str,
        years: int = 2,
        **kwargs,
    ) -> dict[str, list[dict]]:
        """Convenience: fetch the most recent `years` years."""
        end   = datetime.now().year
        start = end - (years - 1)
        return self.fetch(series_ids, start, end, **kwargs)

    def fetch_named(
        self,
        named: dict[str, str],
        start_year: int,
        end_year:   int,
        **kwargs,
    ) -> dict[str, list[dict]]:
        """
        Fetch a labeled dict of series IDs.

        Args:
            named: {"Human label": "SERIES_ID", ...}
        Returns:
            {"Human label": {"data": [...], "catalog": {...}}}
        """
        id_to_label = {v: k for k, v in named.items()}
        raw = self.fetch(list(named.values()), start_year, end_year, **kwargs)
        return {id_to_label[sid]: v for sid, v in raw.items()}

    # ------------------------------------------------------------------
    # Discovery endpoints
    # ------------------------------------------------------------------
    def surveys(self) -> list[dict]:
        """Return all BLS surveys with abbreviation and name."""
        r = self._get(f"{self.BASE_URL}/surveys", timeout=15)
        r.raise_for_status()
        return r.json()["Results"]["survey"]

    def popular(self, survey: str) -> list[str]:
        """Return popular series IDs for a given survey abbreviation."""
        r = self._get(
            f"{self.BASE_URL}/timeseries/popular",
            params={"survey": survey},
            timeout=15,
        )
        r.raise_for_status()
        body = r.json()
        return [
            s["seriesID"]
            for s in body.get("Results", {}).get("series", [])
            if s
        ]

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    @staticmethod
    def latest_obs(series_result: dict) -> dict | None:
        """Return the most recent non-null observation from a fetch result."""
        for obs in series_result.get("data", []):
            if obs["value"] != "-":
                return obs
        return None

    @staticmethod
    def to_rows(results: dict[str, dict]) -> list[dict]:
        """
        Flatten fetch results into a list of dicts suitable for CSV or
        pandas DataFrame ingestion.

        Each row: {series_id, series_title, year, period, period_name, value}
        """
        rows = []
        for sid, s in results.items():
            title = s.get("catalog", {}).get("series_title", "")
            for obs in s.get("data", []):
                if obs["value"] == "-":
                    continue
                rows.append({
                    "series_id":    sid,
                    "series_title": title,
                    "year":         obs["year"],
                    "period":       obs["period"],
                    "period_name":  obs.get("periodName", ""),
                    "value":        obs["value"],
                })
        return rows