"""Typed dataset registry for STRATA acquisition. The registry is the single, version-controlled source of truth describing which datasets STRATA uses, where they legitimately come from, their licensing tier, and which languages they cover. It is intentionally declarative: acquisition code in :mod:`strata.data.acquire` interprets these specs, but no download URLs or credentials live in Python source. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path import yaml # --------------------------------------------------------------------------- # Controlled vocabularies # --------------------------------------------------------------------------- # Licensing / accessibility tiers. These drive default fetch behaviour. TIER_OPEN = "open" # freely downloadable, no gate TIER_GATED_AUTO = "gated_auto" # HF gate with automatic approval on request TIER_GATED_MANUAL = "gated_manual" # HF gate requiring manual approval TIER_LICENSE_REQUIRED = "license_required" # e.g. LDC / signed agreement; no auto DL LICENSE_TIERS = frozenset( {TIER_OPEN, TIER_GATED_AUTO, TIER_GATED_MANUAL, TIER_LICENSE_REQUIRED} ) # Source kinds understood by the acquirer. KIND_HF_SNAPSHOT = "hf_snapshot" # snapshot_download of a HF dataset repo KIND_HF_FILES = "hf_files" # capped/prefix-filtered file selection from a HF repo KIND_HTTP = "http" # direct file/tarball download(s) KIND_GITHUB_ARCHIVE = "github_archive" # pinned GitHub repo tarball KIND_LICENSE_REQUIRED = "license_required" # no download; write acquisition stub SOURCE_KINDS = frozenset( { KIND_HF_SNAPSHOT, KIND_HF_FILES, KIND_HTTP, KIND_GITHUB_ARCHIVE, KIND_LICENSE_REQUIRED, } ) # Pipeline layers, matching the STRATA data stack. "multi" is used in a spec's # languages list to mean "many languages, filtered downstream". LANG_MULTI = "multi" # --------------------------------------------------------------------------- # Spec # --------------------------------------------------------------------------- @dataclass(frozen=True, slots=True) class DatasetSpec: """Declarative description of one dataset.""" id: str title: str layer: str languages: tuple[str, ...] license: str license_tier: str source: dict[str, object] homepage: str = "" notes: str = "" optional: bool = False def __post_init__(self) -> None: self.validate() @property def kind(self) -> str: return str(self.source.get("kind", "")) def covers_language(self, language: str) -> bool: return language in self.languages or LANG_MULTI in self.languages def validate(self) -> None: if not self.id or not self.id.replace("_", "").replace("-", "").isalnum(): raise ValueError( f"dataset id must be non-empty alphanumeric/_/-, got {self.id!r}" ) if self.license_tier not in LICENSE_TIERS: raise ValueError( f"[{self.id}] invalid license_tier {self.license_tier!r}; " f"expected one of {sorted(LICENSE_TIERS)}" ) if not self.languages: raise ValueError(f"[{self.id}] languages must be non-empty") kind = self.kind if kind not in SOURCE_KINDS: raise ValueError( f"[{self.id}] invalid source kind {kind!r}; " f"expected one of {sorted(SOURCE_KINDS)}" ) _validate_source(self.id, kind, self.source) if kind == KIND_LICENSE_REQUIRED and self.license_tier != TIER_LICENSE_REQUIRED: raise ValueError( f"[{self.id}] source kind {KIND_LICENSE_REQUIRED} requires " f"license_tier {TIER_LICENSE_REQUIRED}" ) def _require(dataset_id: str, source: dict[str, object], *keys: str) -> None: for key in keys: if not source.get(key): raise ValueError( f"[{dataset_id}] source kind {source.get('kind')!r} requires " f"non-empty {key!r}" ) def _validate_source(dataset_id: str, kind: str, source: dict[str, object]) -> None: if kind == KIND_HF_SNAPSHOT: _require(dataset_id, source, "repo_id") elif kind == KIND_HF_FILES: _require(dataset_id, source, "repo_id", "prefixes") if not isinstance(source["prefixes"], list): raise ValueError(f"[{dataset_id}] hf_files.prefixes must be a list") elif kind == KIND_HTTP: urls = source.get("urls") if not isinstance(urls, list) or not urls: raise ValueError(f"[{dataset_id}] http.urls must be a non-empty list") for item in urls: if not isinstance(item, dict) or not item.get("url"): raise ValueError(f"[{dataset_id}] each http url needs a 'url' field") elif kind == KIND_GITHUB_ARCHIVE: _require(dataset_id, source, "repo") elif kind == KIND_LICENSE_REQUIRED: _require(dataset_id, source, "acquisition") # --------------------------------------------------------------------------- # Registry # --------------------------------------------------------------------------- @dataclass(frozen=True, slots=True) class Registry: version: int languages_core: tuple[str, ...] datasets: tuple[DatasetSpec, ...] description: str = "" def by_id(self, dataset_id: str) -> DatasetSpec: for spec in self.datasets: if spec.id == dataset_id: return spec raise KeyError(f"unknown dataset id: {dataset_id!r}") def ids(self) -> list[str]: return [spec.id for spec in self.datasets] def layers(self) -> list[str]: seen: dict[str, None] = {} for spec in self.datasets: seen.setdefault(spec.layer, None) return list(seen) def filter( self, *, ids: list[str] | None = None, layers: list[str] | None = None, languages: list[str] | None = None, tiers: list[str] | None = None, include_optional: bool = True, ) -> list[DatasetSpec]: """Return specs matching all provided filters (AND semantics). ``ids`` matching is explicit and always includes optional datasets, so an operator can request an optional dataset by name. """ id_set = set(ids) if ids else None layer_set = set(layers) if layers else None tier_set = set(tiers) if tiers else None selected: list[DatasetSpec] = [] for spec in self.datasets: if id_set is not None and spec.id not in id_set: continue if layer_set is not None and spec.layer not in layer_set: continue if tier_set is not None and spec.license_tier not in tier_set: continue if languages and not any(spec.covers_language(lang) for lang in languages): continue explicitly_named = id_set is not None and spec.id in id_set if spec.optional and not include_optional and not explicitly_named: continue selected.append(spec) return selected def load_registry(path: Path) -> Registry: """Load and validate a registry YAML file.""" if not path.exists(): raise FileNotFoundError(f"registry file not found: {path}") data = yaml.safe_load(path.read_text(encoding="utf-8")) if not isinstance(data, dict): raise ValueError(f"registry {path} must be a mapping at the top level") raw_datasets = data.get("datasets") if not isinstance(raw_datasets, list) or not raw_datasets: raise ValueError(f"registry {path} must contain a non-empty 'datasets' list") specs: list[DatasetSpec] = [] seen_ids: set[str] = set() for raw in raw_datasets: if not isinstance(raw, dict): raise ValueError(f"registry {path}: each dataset entry must be a mapping") spec = _spec_from_dict(raw) if spec.id in seen_ids: raise ValueError(f"registry {path}: duplicate dataset id {spec.id!r}") seen_ids.add(spec.id) specs.append(spec) return Registry( version=int(data.get("version", 1)), languages_core=tuple(data.get("languages_core", []) or []), datasets=tuple(specs), description=str(data.get("description", "")), ) def _spec_from_dict(raw: dict[str, object]) -> DatasetSpec: missing = [key for key in ("id", "layer", "license_tier", "source") if key not in raw] if missing: raise ValueError(f"dataset entry missing required keys: {missing} in {raw!r}") source = raw["source"] if not isinstance(source, dict): raise ValueError(f"[{raw.get('id')}] source must be a mapping") languages = raw.get("languages", []) if not isinstance(languages, list): raise ValueError(f"[{raw.get('id')}] languages must be a list") return DatasetSpec( id=str(raw["id"]), title=str(raw.get("title", raw["id"])), layer=str(raw["layer"]), languages=tuple(str(lang) for lang in languages), license=str(raw.get("license", "unknown")), license_tier=str(raw["license_tier"]), source=dict(source), homepage=str(raw.get("homepage", "")), notes=str(raw.get("notes", "")), optional=bool(raw.get("optional", False)), ) # Backwards-friendly export list. __all__ = [ "DatasetSpec", "Registry", "load_registry", "LICENSE_TIERS", "SOURCE_KINDS", "TIER_OPEN", "TIER_GATED_AUTO", "TIER_GATED_MANUAL", "TIER_LICENSE_REQUIRED", "KIND_HF_SNAPSHOT", "KIND_HF_FILES", "KIND_HTTP", "KIND_GITHUB_ARCHIVE", "KIND_LICENSE_REQUIRED", "LANG_MULTI", ]