From 6dc0664e5c195e829260592918fffe4ce5fca6ff Mon Sep 17 00:00:00 2001 From: Benjamin Thomas Schwertfeger Date: Thu, 13 Nov 2025 08:34:17 +0100 Subject: [PATCH 1/3] Resolve "Add cmethods attributes to output files" --- cmethods/core.py | 48 ++++++++++++++++++++++++++++++++++++++++++-- requirements-dev.txt | 1 + 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/cmethods/core.py b/cmethods/core.py index 714b51b..e757958 100644 --- a/cmethods/core.py +++ b/cmethods/core.py @@ -25,6 +25,7 @@ from __future__ import annotations +from datetime import datetime, timezone from typing import TYPE_CHECKING, Callable, Dict, Optional import xarray as xr @@ -49,6 +50,46 @@ } +def _add_cmethods_metadata( + result: xr.Dataset | xr.DataArray, + method: str, + **kwargs, +) -> xr.Dataset | xr.DataArray: + """ + Add metadata to the result indicating it was processed by python-cmethods. + + :param result: The bias-corrected dataset or dataarray + :param method: The method used for bias correction + :param kwargs: Additional method parameters + :return: Result with added metadata + """ + try: + from importlib.metadata import version # noqa: PLC0415 + + pkg_version = version("python-cmethods") + except Exception: # noqa: BLE001 + pkg_version = "unknown" + + attrs_to_add = { + "cmethods_version": pkg_version, + "cmethods_method": method, + "cmethods_timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC"), + "cmethods_source": "https://github.com/btschwertfeger/python-cmethods", + } + + if kind := kwargs.get("kind"): + attrs_to_add["cmethods_kind"] = kind + if n_quantiles := kwargs.get("n_quantiles"): + attrs_to_add["cmethods_n_quantiles"] = str(n_quantiles) + if group := kwargs.get("group"): + attrs_to_add["cmethods_group"] = str(group) + + if isinstance(result, (xr.Dataset, xr.DataArray)): + result.attrs.update(attrs_to_add) + + return result + + def apply_ufunc( method: str, obs: xr.xarray.core.dataarray.DataArray, @@ -144,6 +185,8 @@ def adjust( :return: The bias corrected/adjusted data set :rtype: xr.xarray.core.dataarray.DataArray | xr.xarray.core.dataarray.Dataset """ + metadata_kwargs = {k: v for k, v in kwargs.items() if k in {"kind", "n_quantiles", "group"}} + kwargs["adjust_called"] = True ensure_xr_dataarray(obs=obs, simh=simh, simp=simp) @@ -159,7 +202,8 @@ def adjust( # mock this function or apply ``CMethods.__apply_ufunc` directly # on your data sets. if kwargs.get("group") is None: - return apply_ufunc(method, obs, simh, simp, **kwargs).to_dataset() + result = apply_ufunc(method, obs, simh, simp, **kwargs).to_dataset() + return _add_cmethods_metadata(result, method, **metadata_kwargs) if method not in SCALING_METHODS: raise ValueError( @@ -204,7 +248,7 @@ def adjust( result = monthly_result if result is None else xr.merge([result, monthly_result]) - return result + return _add_cmethods_metadata(result, method, **metadata_kwargs) __all__ = ["adjust"] diff --git a/requirements-dev.txt b/requirements-dev.txt index dab1c66..c21de7c 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,4 +1,5 @@ dask[distributed] +matplotlib pytest pytest-cov pytest-retry From bd5d79fa6c290ee3de21d5c0e3acd70ff064db33 Mon Sep 17 00:00:00 2001 From: Benjamin Thomas Schwertfeger Date: Thu, 13 Nov 2025 08:42:14 +0100 Subject: [PATCH 2/3] Add tests --- pyproject.toml | 2 ++ tests/test_misc.py | 69 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index a73cf2e..ce9ae49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -264,6 +264,8 @@ min-file-size = 1024 "TID252", # ban relative imports "PTH118", # `os.path.join()` should be replaced by `Path` with `/` operator, "PTH120", # `os.path.dirname()` should be replaced by `Path.parent` + "PLR2004", # magic value in comparison + "PLC2701" # Private name import ] [tool.ruff.lint.flake8-quotes] diff --git a/tests/test_misc.py b/tests/test_misc.py index da8ba92..f7388ca 100644 --- a/tests/test_misc.py +++ b/tests/test_misc.py @@ -15,8 +15,10 @@ import numpy as np import pytest +import xarray as xr from cmethods import adjust +from cmethods.core import _add_cmethods_metadata from cmethods.distribution import ( detrended_quantile_mapping, quantile_delta_mapping, @@ -130,3 +132,70 @@ def test_adjust_failing_no_group_for_distribution(datasets: dict) -> None: n_quantiles=100, group="time.month", ) + + +def test_add_cmethods_metadata_with_dataarray() -> None: + """Test that _add_cmethods_metadata adds correct attributes to a DataArray""" + data = xr.DataArray( + np.array([1, 2, 3, 4, 5]), + dims=["time"], + coords={"time": np.arange(5)}, + ) + + result = _add_cmethods_metadata( + data, + method="linear_scaling", + kind="+", + n_quantiles=100, + group="time.month", + ) + + assert "cmethods_version" in result.attrs + assert "cmethods_method" in result.attrs + assert "cmethods_timestamp" in result.attrs + assert "cmethods_source" in result.attrs + + assert result.attrs["cmethods_method"] == "linear_scaling" + assert result.attrs["cmethods_kind"] == "+" + assert result.attrs["cmethods_n_quantiles"] == "100" + assert result.attrs["cmethods_group"] == "time.month" + assert result.attrs["cmethods_source"] == "https://github.com/btschwertfeger/python-cmethods" + assert "UTC" in result.attrs["cmethods_timestamp"] + + +def test_add_cmethods_metadata_with_dataset() -> None: + """Test that _add_cmethods_metadata adds correct attributes to a Dataset""" + data = xr.Dataset( + { + "temperature": xr.DataArray( + np.array([1, 2, 3, 4, 5]), + dims=["time"], + coords={"time": np.arange(5)}, + ), + }, + ) + + result = _add_cmethods_metadata(data, method="quantile_mapping") + + assert "cmethods_version" in result.attrs + assert "cmethods_method" in result.attrs + assert "cmethods_timestamp" in result.attrs + assert "cmethods_source" in result.attrs + assert result.attrs["cmethods_method"] == "quantile_mapping" + + +def test_add_cmethods_metadata_optional_params() -> None: + """Test that _add_cmethods_metadata handles optional parameters correctly""" + data = xr.DataArray( + np.array([1, 2, 3]), + dims=["time"], + coords={"time": np.arange(3)}, + ) + + result = _add_cmethods_metadata(data, method="variance_scaling") + + assert "cmethods_method" in result.attrs + assert result.attrs["cmethods_method"] == "variance_scaling" + assert "cmethods_kind" not in result.attrs + assert "cmethods_n_quantiles" not in result.attrs + assert "cmethods_group" not in result.attrs From 7263ed713a8058a405b3df43b20d459987f6b753 Mon Sep 17 00:00:00 2001 From: Benjamin Thomas Schwertfeger Date: Thu, 13 Nov 2025 08:44:10 +0100 Subject: [PATCH 3/3] make pre-commit happy --- tests/helper.py | 2 +- tests/test_zarr_dask_compatibility.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/helper.py b/tests/helper.py index 0aacd87..36d380a 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -101,7 +101,7 @@ def get_dataset(data, time, kind: str) -> xr.Dataset: .to_dataset(name=kind) ) - if kind == "+": # noqa: PLR2004 + if kind == "+": some_data = [get_hist_temp_for_lat(val) for val in latitudes] data = np.array( [ diff --git a/tests/test_zarr_dask_compatibility.py b/tests/test_zarr_dask_compatibility.py index 6d5e1bb..d6e827c 100644 --- a/tests/test_zarr_dask_compatibility.py +++ b/tests/test_zarr_dask_compatibility.py @@ -37,7 +37,7 @@ def test_3d_scaling_zarr( kind: str, dask_cluster: Any, # noqa: ARG001 ) -> None: - variable: str = "tas" if kind == "+" else "pr" # noqa: PLR2004 + variable: str = "tas" if kind == "+" else "pr" obsh: xr.DataArray = datasets_from_zarr[kind]["obsh"][variable] obsp: xr.DataArray = datasets_from_zarr[kind]["obsp"][variable] simh: xr.DataArray = datasets_from_zarr[kind]["simh"][variable] @@ -81,7 +81,7 @@ def test_3d_distribution_zarr( kind: str, dask_cluster: Any, # noqa: ARG001 ) -> None: - variable: str = "tas" if kind == "+" else "pr" # noqa: PLR2004 + variable: str = "tas" if kind == "+" else "pr" obsh: XRData_t = datasets_from_zarr[kind]["obsh"][variable] obsp: XRData_t = datasets_from_zarr[kind]["obsp"][variable] simh: XRData_t = datasets_from_zarr[kind]["simh"][variable]