Source code for pamir.catalog

"""Dataset catalog: metadata for all 19 PaMIR datasets."""

import json
from pathlib import Path
from typing import Dict, List, Optional

import pandas as pd

_CATALOG_PATH = Path(__file__).parent / "data" / "catalog.json"
_catalog = None


def _load_raw() -> List[Dict]:
    global _catalog
    if _catalog is None:
        _catalog = json.loads(_CATALOG_PATH.read_text(encoding="utf-8"))
    return _catalog




def load_catalog() -> pd.DataFrame:
    """Return a DataFrame with one row per dataset.

    Columns: id, name, rows, features, defaults, DR, source, license,
    geography, product, target_definition.
    """
    return pd.DataFrame(_load_raw()).set_index("id")





def list_datasets() -> List[str]:
    """Return sorted list of dataset ids."""
    return sorted(e["id"] for e in _load_raw())



# Download sources that require credentials to fetch.
_CREDENTIAL_KINDS = {"kaggle", "kaggle_competition"}


def _needs_credentials(entry: Dict) -> bool:
    return entry.get("download", {}).get("kind") in _CREDENTIAL_KINDS


def open_datasets() -> List[str]:
    """Ids whose source needs **no credentials** to download.

    These come from UCI / GitHub / HuggingFace and work on a
    fresh install with no setup; the rest are on Kaggle and need Kaggle
    credentials (see ``pamir.download``).
    """
    return sorted(e["id"] for e in _load_raw() if not _needs_credentials(e))




def dataset_info(dataset_id: str) -> Dict:
    """Return full metadata dict for one dataset.

    Raises KeyError if the id is not in the catalog.
    """
    for e in _load_raw():
        if e["id"] == dataset_id:
            d = dict(e)
            d["needs_credentials"] = _needs_credentials(e)
            return d
    raise KeyError(
        f"Unknown dataset '{dataset_id}'. "
        f"Available: {', '.join(list_datasets())}"
    )