stack/tests/lead/conftest.py

"""Backends for the lead Expr tests: every frame test runs on polars and on
DuckDB through narwhals."""

from __future__ import annotations

from collections.abc import Iterator
from typing import Any

import duckdb
import narwhals as nw
import polars as pl
import pytest


class Backend:
    """Build a native frame on one backend and read results back as rows."""

    def __init__(self, name: str) -> None:
        self.name = name
        self.con = duckdb.connect() if name == "duckdb" else None

    def frame(self, data: Any, schema: dict[str, Any] | None = None) -> Any:
        df = pl.DataFrame(data, schema=schema)
        return df if self.con is None else self.con.from_arrow(df.to_arrow())

    def rows(self, native: Any) -> list[dict[str, Any]]:
        return nw.from_native(native).lazy().collect("polars").to_native().to_dicts()

    def select(self, data: Any, schema: dict[str, Any] | None = None, **exprs: nw.Expr):
        frame = nw.from_native(self.frame(data, schema))
        return self.rows(frame.select(**exprs).to_native())

    def close(self) -> None:
        if self.con is not None:
            self.con.close()


@pytest.fixture(params=["polars", "duckdb"])
def backend(request: pytest.FixtureRequest) -> Iterator[Backend]:
    engine = Backend(request.param)
    try:
        yield engine
    finally:
        engine.close()