py_runner.py

import asyncio
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, List

ROOT = Path(__file__).resolve().parent
REPO_ROOT = ROOT.parent
PY_CLIENT_PATH = REPO_ROOT / "dclimate-client-py"

# Use local clone of the Python client
sys.path.insert(0, str(PY_CLIENT_PATH))

from dclimate_client_py import dClimateClient  # type: ignore  # pylint: disable=wrong-import-position


def parse_case() -> dict:
    raw = sys.stdin.read()
    try:
        return json.loads(raw)
    except json.JSONDecodeError as exc:
        raise SystemExit(f"Invalid case payload: {exc}") from exc


def parse_iso8601(value: str) -> datetime:
    if value.endswith("Z"):
        value = value.replace("Z", "+00:00")
    dt = datetime.fromisoformat(value)
    # Convert to UTC and drop tzinfo to match tz-naive dataset coords
    if dt.tzinfo is not None:
        dt = dt.astimezone(timezone.utc).replace(tzinfo=None)
    return dt


def flatten(values: Any) -> List[Any]:
    flattened: List[Any] = []
    stack = [values]
    while stack:
        current = stack.pop()
        if isinstance(current, list):
            stack.extend(reversed(current))
        elif current is None:
            flattened.append(None)
        elif isinstance(current, (int, float)):
            flattened.append(float(current))
    return flattened


def sanitize_metadata(metadata: dict) -> dict:
    safe = {k: v for k, v in metadata.items() if k not in {"slug"}}
    if "timestamp" in safe and safe["timestamp"] is not None:
        safe["timestamp"] = int(safe["timestamp"])
    return safe


async def run_case(test_case: dict) -> dict:
    python_opts = test_case.get("python", {})
    selection = test_case.get("selection", {}) or {}
    request = test_case.get("request", {}) or {}
    variable = test_case.get("variable") or request.get("dataset")

    async with dClimateClient() as client:
        dataset, metadata = await client.load_dataset(
            dataset=request.get("dataset"),
            collection=request.get("collection"),
            variant=request.get("variant"),
            cid=python_opts.get("cid"),
            return_xarray=False,
        )

        selected_variable = None
        # Ensure dataset is scoped to a data variable if needed
        if hasattr(dataset, "data"):
            data_vars = list(dataset.data.data_vars)
            if variable in data_vars:
                dataset = dataset.use(variable)
                selected_variable = variable
            elif len(data_vars) == 1:
                dataset = dataset.use(data_vars[0])
                selected_variable = data_vars[0]
            elif len(data_vars) > 1:
                dataset = dataset.use(data_vars[0])
                selected_variable = data_vars[0]

        if selection.get("point"):
            point = selection["point"]
            dataset = dataset.point(
                latitude=point["latitude"],
                longitude=point["longitude"],
                snap_to_grid=point.get("snap_to_grid", True),
            )

        if selection.get("timeRange"):
            time_range = selection["timeRange"]
            dataset = dataset.time_range(
                parse_iso8601(time_range["start"]),
                parse_iso8601(time_range["end"]),
            )

        payload = dataset.as_dict()
        values = flatten(payload.get("data", []))
        return {
            "origin": "python",
            "metadata": sanitize_metadata(metadata),
            "selectedVariable": selected_variable,
            "variables": list(dataset.data.data_vars),
            "coords": {
                "dimensions": {k: int(v) for k, v in dataset.data.sizes.items()},
                "coordKeys": list(dataset.data.coords.keys()),
            },
            "values": values,
        }


async def main():
    test_case = parse_case()
    result = await run_case(test_case)
    sys.stdout.write(json.dumps(result, default=str))


if __name__ == "__main__":
    asyncio.run(main())