Resolve "Add cmethods attributes to output files" (#177)

Author: btschwertfeger

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"]

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]

requirements-dev.txt

@@ -1,4 +1,5 @@
 dask[distributed]
+matplotlib
 pytest
 pytest-cov
 pytest-retry

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(
             [

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

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]