sdmxlib 0.8.0__py3-none-any.whl

sdmxlib/api/query.py

@@ -0,0 +1,428 @@
+"""GetQuery[T], Query[T], and DataQuery — artefact and data queries."""
+
+from collections.abc import Callable
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Self
+
+import polars as pl
+from attrs import frozen
+
+from sdmxlib.api.filters import DimFilter
+
+if TYPE_CHECKING:
+    from sdmxlib.model.dataset import Dataset
+
+
+@frozen
+class _KeySpec:
+    """Bundled key inputs passed from DataQuery to registry closures."""
+
+    raw_key: str = "all"
+    slices: tuple[DimFilter, ...] = ()
+    measures: tuple[str, ...] | None = None
+    measure_dim_hint: str | None = None
+
+
+def _to_references(reference_type: type) -> str:
+    """Map a model type to its SDMX REST ``references=`` parameter string."""
+    sdmx_class: str = getattr(reference_type, "sdmx_class", reference_type.__name__)
+    return sdmx_class.lower()
+
+
+class GetQuery[T]:
+    """Single-artefact lookup with configurable resolution depth.
+
+    Returned by ``reg.get()``. Call ``.fetch()`` for an unresolved skeleton
+    or ``.resolve()`` to follow references.
+
+    Example:
+        reg.get(sl.DataStructure, agency="ECB", id="ECB_EXR1").fetch()
+        reg.get(sl.DataStructure, agency="ECB", id="ECB_EXR1").resolve()
+        reg.get(sl.DataStructure, agency="ECB", id="ECB_EXR1").resolve(sl.Codelist)
+        reg.get(sl.DataStructure, agency="ECB", id="ECB_EXR1").resolve(partial=False)
+    """
+
+    def __init__(
+        self,
+        fetcher: Callable[[], T],
+        resolver: Callable[[str, bool], T],
+    ) -> None:
+        self._fetcher = fetcher
+        self._resolver = resolver
+
+    def fetch(self) -> T:
+        """Fetch the artefact skeleton only (``detail=full``, ``references=none``)."""
+        return self._fetcher()
+
+    def resolve(
+        self,
+        reference_type: type | None = None,
+        *,
+        partial: bool = True,
+    ) -> T:
+        """Fetch the artefact with references resolved.
+
+        Args:
+            reference_type: If given, resolve only Refs of this type.
+                            If ``None`` (default), resolve all references.
+            partial: If ``True`` (default), referenced ItemSchemes are trimmed
+                     to items allowed by ContentConstraints
+                     (``detail=referencepartial``). Produces smaller, more
+                     useful responses for the common case. Set to ``False``
+                     for full metadata audits (``detail=full``).
+
+        Wire mapping:
+            resolve()                   → detail=referencepartial, references=descendants
+            resolve(partial=False)      → detail=full,             references=descendants
+            resolve(sl.Codelist)        → detail=referencepartial, references=codelist
+            resolve(sl.Codelist, partial=False) → detail=full,     references=codelist
+        """
+        ref_str = "descendants" if reference_type is None else _to_references(reference_type)
+        return self._resolver(ref_str, partial)
+
+
+class Query[T]:
+    """Lazy collection search. Fetches on terminal call.
+
+    Server-side constraint (agency) is set at construction.
+    Client-side refinements (filter, sort_by, limit) chain before the terminal.
+
+    Terminals:
+        browse()          — identity stubs only (no items populated)
+        fetch()           — full artefacts, Refs unresolved
+        resolve()         — full artefacts, Refs resolved (constrained by default)
+        resolve(partial=False) — full artefacts, Refs resolved, all codes
+
+    Example:
+        reg.query(sl.Codelist, agency="ESTAT").fetch()
+
+        (
+            reg.query(sl.Dataflow, agency="ECB")
+            .filter(lambda df: "EXR" in df.id)
+            .sort_by("id")
+            .limit(5)
+            .resolve()
+        )
+
+        reg.query(sl.Dataflow, agency="ECB").resolve(sl.DataStructure)
+        reg.query(sl.Dataflow, agency="ECB").browse()
+    """
+
+    def __init__(
+        self,
+        fetcher: Callable[[], list[T]],
+        resolver: Callable[[str, bool], list[T]],
+        browser: Callable[[], list[T]],
+    ) -> None:
+        self._fetcher = fetcher
+        self._resolver = resolver
+        self._browser = browser
+        self._filters: list[Callable[[T], bool]] = []
+        self._order_key: str | Callable[[T], Any] | None = None
+        self._descending: bool = False
+        self._limit: int | None = None
+
+    # ── Refinement (return Self) ──────────────────────────────────────────────
+
+    def filter(self, predicate: Callable[[T], bool]) -> Self:
+        """Add a client-side filter predicate."""
+        self._filters.append(predicate)
+        return self
+
+    def sort_by(self, key: str | Callable[[T], Any], *, descending: bool = False) -> Self:
+        """Sort results by attribute name or key function."""
+        self._order_key = key
+        self._descending = descending
+        return self
+
+    def limit(self, n: int) -> Self:
+        """Truncate results to at most n items."""
+        self._limit = n
+        return self
+
+    # ── Internal pipeline ─────────────────────────────────────────────────────
+
+    def _apply(self, items: list[T]) -> list[T]:
+        for pred in self._filters:
+            items = [item for item in items if pred(item)]
+        if self._order_key is not None:
+            order_key = self._order_key
+            key_fn: Callable[[T], Any] = order_key if callable(order_key) else lambda item: getattr(item, order_key)
+            items = sorted(items, key=key_fn, reverse=self._descending)
+        if self._limit is not None:
+            items = items[: self._limit]
+        return items
+
+    # ── Terminals ─────────────────────────────────────────────────────────────
+
+    def browse(self) -> list[T]:
+        """Fetch lightweight identity stubs (``detail=allstubs``, ``references=none``).
+
+        Each artefact has identity fields (id, agency, version, name) but
+        item collections are empty. Use for catalogue scanning before deciding
+        which artefacts to fetch in full.
+
+        Example:
+            for df in reg.query(sl.Dataflow, agency="ECB").browse():
+                print(df.id, df.name.get("en"))
+        """
+        return self._apply(self._browser())
+
+    def fetch(self) -> list[T]:
+        """Fetch artefacts with unresolved Refs (``detail=full``, ``references=none``)."""
+        return self._apply(self._fetcher())
+
+    def fetch_one(self) -> T:
+        """Fetch exactly one artefact. Raises if count != 1."""
+        results = self.fetch()
+        if len(results) != 1:
+            msg = f"Expected exactly 1 result, got {len(results)}"
+            raise ValueError(msg)
+        return results[0]
+
+    def fetch_first(self) -> T:
+        """Fetch the first matching artefact. Raises if none found."""
+        results = self.fetch()
+        if not results:
+            msg = "No results found"
+            raise ValueError(msg)
+        return results[0]
+
+    def fetch_optional(self) -> T | None:
+        """Fetch the first matching artefact, or None if none found."""
+        results = self.fetch()
+        return results[0] if results else None
+
+    def resolve(
+        self,
+        reference_type: type | None = None,
+        *,
+        partial: bool = True,
+    ) -> list[T]:
+        """Fetch artefacts with references resolved.
+
+        Args:
+            reference_type: If given, resolve only Refs of this type.
+            partial: Controls whether referenced ItemSchemes are trimmed by
+                     ContentConstraints. ``True`` (default) is faster and
+                     more useful for data queries; ``False`` for full audits.
+        """
+        ref_str = "descendants" if reference_type is None else _to_references(reference_type)
+        return self._apply(self._resolver(ref_str, partial))
+
+    def resolve_one(
+        self,
+        reference_type: type | None = None,
+        *,
+        partial: bool = True,
+    ) -> T:
+        """Resolve exactly one artefact. Raises if count != 1."""
+        results = self.resolve(reference_type, partial=partial)
+        if len(results) != 1:
+            msg = f"Expected exactly 1 result, got {len(results)}"
+            raise ValueError(msg)
+        return results[0]
+
+    def resolve_first(
+        self,
+        reference_type: type | None = None,
+        *,
+        partial: bool = True,
+    ) -> T:
+        """Resolve the first matching artefact. Raises if none found."""
+        results = self.resolve(reference_type, partial=partial)
+        if not results:
+            msg = "No results found"
+            raise ValueError(msg)
+        return results[0]
+
+    def resolve_optional(
+        self,
+        reference_type: type | None = None,
+        *,
+        partial: bool = True,
+    ) -> T | None:
+        """Resolve the first matching artefact, or None if none found."""
+        results = self.resolve(reference_type, partial=partial)
+        return results[0] if results else None
+
+    def count(self) -> int:
+        """Return the number of matching artefacts."""
+        return len(self.fetch())
+
+
+class DataQuery:
+    """Lazy data fetch — resolves to a :class:`~sdmxlib.model.dataset.Dataset`.
+
+    Returned by :meth:`~sdmxlib.api.registry.RestRegistry.data`. Use
+    :meth:`fetch` for in-memory loading or :meth:`fetch_to` to stream
+    directly to a file for large datasets. Use :meth:`slice` and
+    :meth:`measures` to filter what comes over the wire.
+
+    Example:
+        # Raw positional key string (backwards compatible)
+        ds = reg.data(flow, key="M.USD.EUR.SP00.A").fetch()
+
+        # Named dimension filter — no positional knowledge needed
+        ds = (
+            reg.data(flow)
+            .slice(sl.dim("FREQ") == "M", sl.dim("CURRENCY") == "USD")
+            .fetch()
+        )
+
+        # Measure selection + pivot for long-format SDMX 2.1 flows
+        df = (
+            reg.data(flow)
+            .slice(sl.dim("REF_AREA").isin(["AT", "DE"]))
+            .measures(["EMP", "UNE"])
+            .pivot()
+            .collect()
+        )
+
+        # File-backed streaming for large data
+        ds = reg.data(sl.Dataflow, agency="ESTAT", id="LFS").fetch_to("lfs.csv")
+        ds.sink_parquet("lfs.parquet")
+    """
+
+    def __init__(
+        self,
+        fetcher: "Callable[[_KeySpec], Dataset]",
+        streamer: "Callable[[Path, _KeySpec], Dataset]",
+        raw_key: str = "all",
+        measure_dim: str | None = None,
+    ) -> None:
+        self._fetcher = fetcher
+        self._streamer = streamer
+        self._raw_key = raw_key
+        self._measure_dim = measure_dim
+        self._slices: list[DimFilter] = []
+        self._measures: list[str] | None = None
+
+    def slice(self, *filters: DimFilter) -> Self:
+        """Add server-side dimension filters to the SDMX key.
+
+        Each filter is a :class:`~sdmxlib.api.filters.DimFilter` produced by
+        :func:`~sdmxlib.dim`. Calling ``.slice()`` multiple times accumulates
+        filters — each call is additive.
+
+        Unspecified dimensions become wildcards. Raises ``ValueError`` if
+        ``key=`` was also passed to
+        :meth:`~sdmxlib.api.registry.RestRegistry.data`.
+
+        Example:
+            import sdmxlib as sl
+
+            reg.data(flow).slice(sl.dim("FREQ") == "M").fetch()
+            reg.data(flow).slice(
+                sl.dim("FREQ") == "M",
+                sl.dim("REF_AREA").isin(["AT", "DE"]),
+            ).fetch()
+
+            # Chained — equivalent to a single call
+            (
+                reg.data(flow)
+                .slice(sl.dim("FREQ") == "M")
+                .slice(sl.dim("REF_AREA") == "AT")
+                .fetch()
+            )
+        """
+        if self._raw_key != "all":
+            msg = "Cannot combine .slice() with key= — use one or the other"
+            raise ValueError(msg)
+        self._slices.extend(filters)
+        return self
+
+    def measures(self, codes: "str | list[str]") -> Self:
+        """Select which measures or indicators to fetch.
+
+        For SDMX 2.1 flows with a ``MeasureDimension`` (e.g. ``INDIC_EM``),
+        this filters the indicator key segment — controlling which columns
+        appear after :meth:`pivot`. For SDMX 3.0 multi-measure flows this
+        will become the ``measures=`` query parameter.
+
+        ``codes`` is a single code string or a list of code strings.
+
+        Example:
+            # Fetch only EMP and UNE indicators, then pivot to wide
+            reg.data(flow).measures(["EMP", "UNE"]).pivot().collect()
+
+            # Single measure
+            reg.data(flow).measures("EMP").fetch()
+        """
+        self._measures = [codes] if isinstance(codes, str) else list(codes)
+        return self
+
+    def _spec(self) -> _KeySpec:
+        return _KeySpec(
+            raw_key=self._raw_key,
+            slices=tuple(self._slices),
+            measures=tuple(self._measures) if self._measures is not None else None,
+            measure_dim_hint=self._measure_dim,
+        )
+
+    def fetch(self) -> "Dataset":
+        """Fetch data into memory and return a Dataset.
+
+        The full response is loaded as bytes then wrapped in a ``pl.LazyFrame``.
+        Suitable for datasets that fit comfortably in memory.
+        """
+        return self._fetcher(self._spec())
+
+    def pivot(self) -> pl.LazyFrame:
+        """Fetch data and pivot to wide format in one step.
+
+        Equivalent to ``reg.data(...).fetch().pivot()``. Requires a
+        ``MeasureDimension`` in the DSD — raises ``ValueError`` otherwise.
+
+        Returns a ``pl.LazyFrame`` — call ``.collect()`` to materialise or
+        ``.sink_parquet()`` to stream directly to a file.
+
+        Example:
+            wide = reg.data(flow).measures(["EMP", "UNE"]).pivot().collect()
+            reg.data(flow).pivot().sink_parquet("wide.parquet")
+        """
+        return self._fetcher(self._spec()).pivot()
+
+    def pivot_to(self, path: "str | Path") -> None:
+        """Stream data to a file and sink the pivoted wide table.
+
+        Equivalent to ``fetch_to(path).pivot().sink_parquet(...)`` but uses
+        a single temp file internally. The long-format file is written by
+        httpx chunked streaming; the pivot happens lazily when sinking.
+
+        Args:
+            path: Destination Parquet file path.
+
+        Example:
+            reg.data(flow).measures(["EMP", "UNE"]).pivot_to("wide.parquet")
+        """
+        import tempfile  # noqa: PLC0415
+
+        suffix = ".csv"
+        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
+            tmp_path = Path(tmp.name)
+        try:
+            ds = self._streamer(tmp_path, self._spec())
+            ds.pivot().sink_parquet(Path(path))
+        finally:
+            tmp_path.unlink(missing_ok=True)
+
+    def fetch_to(self, path: "str | Path") -> "Dataset":
+        """Stream data directly to a file and return a file-backed Dataset.
+
+        Uses httpx chunked streaming — the response is never fully loaded into
+        memory. The returned ``Dataset.data`` is a genuine ``pl.LazyFrame``
+        backed by ``pl.scan_csv``, so :meth:`~sdmxlib.model.dataset.Dataset.sink_parquet`
+        and :meth:`~sdmxlib.model.dataset.Dataset.sink_csv` are truly streaming.
+
+        The caller owns the file — it is not deleted when the Dataset is
+        garbage collected.
+
+        Args:
+            path: Destination file path.
+
+        Example:
+            ds = reg.data(sl.Dataflow, agency="ESTAT", id="LFS").fetch_to("lfs.csv")
+            ds.sink_parquet("lfs.parquet")
+        """
+        return self._streamer(Path(path), self._spec())