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


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):
        self.api_key = api_key

    # ------------------------------------------------------------------
    # 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,
            }
            r = requests.post(
                f"{self.BASE_URL}/timeseries/data/",
                json=payload,
                timeout=30,
            )
            r.raise_for_status()
            body = r.json()
            if body["status"] != "REQUEST_SUCCEEDED":
                raise RuntimeError(f"BLS API error: {body['message']}")
            for s in body["Results"]["series"]:
                results[s["seriesID"]] = {
                    "data":    s["data"],
                    "catalog": s.get("catalog", {}),
                }

        return results

    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 = requests.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 = requests.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