Source code for fews_py_wrapper.utils

import inspect
import io
import zipfile
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from tempfile import TemporaryDirectory
from typing import Any, Callable

import xarray as xr

__all__ = [
    "format_datetime",
    "convert_netcdf_zip_response_to_xarray",
    "format_time_args",
    "get_function_arg_names",
]




def format_datetime(dt: datetime, time_format: str = "%Y-%m-%dT%H:%M:%SZ") -> str:
    """Format a datetime object to a string suitable for FEWS web services."""
    if not dt.tzinfo:
        raise ValueError("Datetime object must be timezone-aware.")
    dt = dt.astimezone(timezone.utc)
    return dt.strftime(time_format)





def convert_netcdf_zip_response_to_xarray(response_content: bytes) -> list[xr.Dataset]:
    """Convert FEWS NetCDF content to xarray datasets.

    ZIP responses are returned as one loaded dataset per NetCDF member, in the
    same order as the ZIP archive. The original ZIP member filename without the
    file extension is preserved in each dataset's ``fews_zip_member_filename``
    attribute. Single-file NetCDF responses are returned as a one-item list.
    """
    try:
        datasets = _load_netcdf_member_datasets(response_content)
    except zipfile.BadZipFile:
        try:
            datasets = [_load_single_netcdf_dataset(response_content)]
        except Exception as exc:
            raise ValueError(
                "Expected FEWS PI_NETCDF content as a ZIP archive containing NetCDF "
                "files or as a single NetCDF file."
            ) from exc

    return datasets



def _load_netcdf_member_datasets(response_content: bytes) -> list[xr.Dataset]:
    """Load each NetCDF member from a FEWS ZIP response."""
    with zipfile.ZipFile(io.BytesIO(response_content)) as zip_file:
        netcdf_members = [
            member
            for member in zip_file.infolist()
            if not member.is_dir()
            and member.filename.lower().endswith((".nc", ".nc4", ".cdf"))
        ]
        if not netcdf_members:
            raise ValueError("ZIP response did not contain any .nc files.")

        datasets: list[xr.Dataset] = []
        with TemporaryDirectory() as temp_dir:
            for index, member in enumerate(netcdf_members):
                extracted_path = _write_zip_member_to_temp_path(
                    zip_file, member, Path(temp_dir), index
                )
                with xr.open_dataset(extracted_path) as dataset:
                    loaded_dataset = dataset.load()
                    loaded_dataset.attrs["fews_zip_member_filename"] = (
                        _strip_zip_member_suffix(member.filename)
                    )
                    datasets.append(loaded_dataset)
        return datasets


def _load_single_netcdf_dataset(response_content: bytes) -> xr.Dataset:
    """Load a single-file NetCDF response."""
    try:
        with TemporaryDirectory() as temp_dir:
            netcdf_path = Path(temp_dir) / "response.nc"
            netcdf_path.write_bytes(response_content)
            with xr.open_dataset(netcdf_path) as dataset:
                return dataset.load()
    except Exception as exc:
        raise ValueError(
            "Expected FEWS PI_NETCDF content as a ZIP archive containing NetCDF "
            "files or as a single NetCDF file."
        ) from exc


def _strip_zip_member_suffix(filename: str) -> str:
    """Return a ZIP member filename without its final file extension."""
    return PurePosixPath(filename).with_suffix("").as_posix()


def _write_zip_member_to_temp_path(
    zip_file: zipfile.ZipFile,
    member: zipfile.ZipInfo,
    temp_dir: Path,
    index: int,
) -> Path:
    """Write a ZIP member to a controlled path under the temp directory."""
    extracted_path = temp_dir / f"member_{index}.nc"
    with zip_file.open(member) as src, extracted_path.open("wb") as dst:
        dst.write(src.read())
    return extracted_path




def format_time_args(*args: None | datetime) -> list[None | str]:
    """Format a list of datetime arguments to strings suitable for web services."""
    formatted_args: list[str | None] = []
    for dt in args:
        if dt is None:
            formatted_args.append(None)
        else:
            formatted_args.append(format_datetime(dt))
    return formatted_args





def get_function_arg_names(func: Callable[..., Any]) -> list[str]:
    """Get the argument names of a function."""
    return list(inspect.signature(func).parameters)