tsfm_bench/data/interface.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterator, Protocol, runtime_checkable

import numpy as np


@dataclass(frozen=True)
class DatasetMetadata:
    """Static metadata for a benchmark dataset configuration."""

    name: str
    domain: str
    frequency: str
    num_variates: int
    terms: tuple[str, ...] = ("short",)
    source_type: str = "synthetic"


@dataclass(frozen=True)
class TimeSeriesRecord:
    """One multivariate or univariate time series instance."""

    item_id: str
    target: np.ndarray
    start: str
    freq: str


@runtime_checkable
class TimeSeriesDataSource(Protocol):
    """Streaming data source for benchmark datasets.

    Real-world backends (HF datasets, databases, APIs) should implement this
    protocol and yield ``TimeSeriesRecord`` instances lazily via ``stream()``.
    """

    def list_datasets(self) -> list[str]:
        ...

    def get_metadata(self, name: str) -> DatasetMetadata:
        ...

    def stream(self, name: str) -> Iterator[TimeSeriesRecord]:
        ...