feat: add augment_from_xarray for EOProduct (#113)

Co-authored-by: Sylvain Brunato <sylvain.brunato@c-s.fr>

Author: timothey-cs

eodag_cube/api/product/_product.py

@@ -46,6 +46,7 @@
 from eodag_cube.api.product._assets import AssetsDict
 from eodag_cube.types import XarrayDict
 from eodag_cube.utils.exceptions import DatasetCreationError
+from eodag_cube.utils.metadata import build_bands, build_cube_metadata, merge_bands
 from eodag_cube.utils.xarray import try_open_dataset
 
 logger = logging.getLogger("eodag-cube.api.product")
@@ -348,3 +349,61 @@ def to_xarray(
             xd.sort()
 
             return xd
+
+    def augment_from_xarray(
+        self,
+        roles: Iterable[str] = {"data", "data-mask"},
+    ) -> EOProduct:
+        """
+        Annotate the product properties and assets with STAC metadata got by fetching its xarray representation.
+
+        :param roles: (optional) roles of assets that must be fetched
+        :returns: updated EOProduct
+        """
+        if not self.assets:
+            try:
+                xd = self.to_xarray(roles=roles)
+            except Exception:
+                return self
+
+            dimensions, variables, proj_info = build_cube_metadata(xd)
+            self.properties["cube:dimensions"] = dimensions
+            self.properties["cube:variables"] = variables
+            self.properties["bands"] = build_bands(xd)
+            for key, value in proj_info.items():
+                self.properties[key] = value
+
+        else:
+            # have roles been set in assets ?
+            roles_exist = any("roles" in a for a in self.assets.values())
+            for asset_key, asset in self.assets.items():
+                try:
+                    asset_roles = asset.get("roles", [])
+                    if (
+                        roles
+                        and asset_roles
+                        and not any(r in asset_roles for r in roles)
+                        or not roles
+                        or not roles_exist
+                    ):
+                        continue
+                    xd = self.to_xarray(asset_key=asset_key, roles=roles)
+                except Exception:
+                    continue
+
+                dimensions, variables, proj_info = build_cube_metadata(xd)
+                asset["cube:dimensions"] = dimensions
+                asset["cube:variables"] = variables
+                for key, value in proj_info.items():
+                    asset[key] = value
+
+                has_band_data = any("band_data" in ds.data_vars for ds in xd.values())
+
+                if has_band_data:
+                    generated_bands = build_bands(xd)
+                    if "bands" in asset:
+                        asset["bands"] = merge_bands(asset["bands"], generated_bands)
+                    else:
+                        asset["bands"] = generated_bands
+
+        return self

eodag_cube/utils/metadata.py

@@ -0,0 +1,248 @@
+# -*- coding: utf-8 -*-
+# Copyright 2026, CS GROUP - France, http://www.c-s.fr
+#
+# This file is part of EODAG project
+#     https://www.github.com/CS-SI/EODAG
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Metadata-related utilities for eodag-cube."""
+
+from math import isnan
+from typing import Any, Union
+
+import numpy as np
+from xarray import DataArray, Dataset
+
+from eodag_cube.types import XarrayDict
+
+
+def extract_projection_info(ds: Dataset) -> dict[str, Any]:
+    """
+    Extract projection information from a :class:`xarray.Dataset`.
+
+    :param ds: :class:`xarray.Dataset` to extract projection information from
+    :return: dictionary with projection information
+    """
+    proj_info: dict[str, Any] = {}
+
+    epsg_code = 4326
+    proj_bbox = None
+
+    if hasattr(ds, "rio") and ds.rio.crs is not None:
+        epsg_code = ds.rio.crs.to_epsg() or 4326
+        try:
+            proj_bbox = list(ds.rio.bounds())
+        except Exception:
+            proj_bbox = None
+
+    proj_info["proj:code"] = f"EPSG:{epsg_code}"
+    if proj_bbox is not None:
+        proj_info["proj:bbox"] = proj_bbox
+    proj_info["proj:shape"] = list(ds.sizes.values())
+    return proj_info
+
+
+def _get_nodata_value(var: DataArray) -> Union[float, str, None]:
+    """
+    Get nodata value from a variable's attributes or return a default value.
+
+    :param var: variable to get nodata value from
+    :return: nodata value
+    """
+    if "nodata" in var.attrs:
+        value = var.attrs["nodata"]
+    elif "_FillValue" in var.encoding:
+        value = var.encoding["_FillValue"]
+    elif "missing_value" in var.encoding:
+        value = var.encoding["missing_value"]
+    elif hasattr(var, "rio"):
+        value = getattr(var.rio, "encoded_nodata", None)
+        if value is None:
+            value = getattr(var.rio, "nodata", None)
+    else:
+        return None
+
+    if value is None:
+        return None
+
+    # handle NaN
+    value = float(value)
+    if isnan(value):
+        return str(value)
+
+    return value
+
+
+def set_variables(ds: Dataset) -> dict[str, Any]:
+    """
+    Set variables metadata from a :class:`xarray.Dataset`.
+
+    :param ds: :class:`xarray.Dataset` to extract variables metadata from
+    :return: dictionary with variables metadata
+    """
+    variables: dict[str, dict] = {}
+    auxiliary_geo_vars: dict[str, str] = {
+        "latitude": "Latitude",
+        "longitude": "Longitude",
+    }
+    for var_name, var in ds.data_vars.items():
+        variables[str(var_name)] = {
+            "dimensions": list(var.dims),
+            "type": "data",
+            "data_type": str(var.dtype),
+        }
+        if desc := var.attrs.get("description"):
+            variables[str(var_name)]["description"] = desc
+        variables[str(var_name)]["nodata"] = _get_nodata_value(var)
+
+    for aux_name, desc in auxiliary_geo_vars.items():
+        if aux_name in ds:
+            var = ds[aux_name]
+
+            if aux_name in variables:
+                continue
+            if aux_name in ds.dims:
+                continue
+
+            variables[aux_name] = {
+                "dimensions": list(var.dims),
+                "type": "auxiliary",
+                "description": desc,
+                "data_type": str(var.dtype),
+            }
+            variables[aux_name]["nodata"] = _get_nodata_value(var)
+
+    return variables
+
+
+def build_cube_metadata(ds_dict: XarrayDict) -> tuple[dict, dict, dict]:
+    """
+    Build datacube and projection metadata from a dict of :class:`xarray.Dataset`.
+
+    :param ds_dict: input xarray dict
+    :return: tuple of 3 dicts for cube dimensions, cube variables and projection info
+    """
+    dimensions: dict[str, dict] = {}
+    variables: dict[str, dict] = {}
+
+    for ds in ds_dict.values():
+        proj_info: dict[str, Any] = extract_projection_info(ds)
+
+        # Dimensions
+        for dim_name in ds.sizes.keys():
+            dim_name_str = str(dim_name)
+
+            # Type
+            dim_type = (
+                "spatial"
+                if dim_name_str in ("x", "y", "lon", "lat")
+                else "temporal"
+                if dim_name_str == "time"
+                else "other"
+            )
+
+            dim_entry: dict[str, Any] = {"type": dim_type}
+
+            if dim_type == "spatial":
+                # Axis
+                if dim_name_str in ("x", "lon"):
+                    dim_entry["axis"] = "x"
+                elif dim_name_str in ("y", "lat"):
+                    dim_entry["axis"] = "y"
+                elif dim_name_str == "z":
+                    dim_entry["axis"] = "z"
+
+                proj_code = proj_info.get("proj:code", "EPSG:4326")
+                try:
+                    dim_entry["reference_system"] = int(proj_code.split(":")[-1])
+                except ValueError:
+                    pass
+
+            if dim_name_str in ds.coords:
+                values = ds[dim_name_str].values
+                if values.ndim == 1:
+                    if values.size <= 10:
+                        dim_entry["values"] = values.tolist()
+                    else:
+                        dim_entry["extent"] = (
+                            [float(values.min()), float(values.max())]
+                            if np.issubdtype(values.dtype, np.number)
+                            else [str(values.min()), str(values.max())]
+                        )
+                        diffs = np.diff(values)
+                        if np.allclose(diffs, diffs[0]):
+                            dim_entry["step"] = (
+                                float(diffs[0]) if np.issubdtype(values.dtype, np.number) else str(diffs[0])
+                            )
+                else:
+                    dim_entry["extent"] = [float(np.nanmin(values)), float(np.nanmax(values))]
+
+            dimensions[dim_name_str] = dim_entry
+
+        # Variables
+        var_ds = set_variables(ds)
+        variables.update(var_ds)
+
+    return dimensions, variables, proj_info
+
+
+def build_bands(xd: XarrayDict) -> list[dict]:
+    """
+    Build STAC bands metadata from xarray datasets.
+
+    If names are not available, use generic band names.
+
+    :param xd: input xarray dict
+    :return: list of bands metadata
+    """
+    band_count = 0
+
+    for ds in xd.values():
+        for var in ds.data_vars.values():
+            for dim in var.dims:
+                if str(dim).lower() in ("band", "bands"):
+                    band_count = ds.sizes[dim]
+                    break
+            if band_count:
+                break
+
+        if band_count:
+            break
+
+    if band_count == 0:
+        band_count = len(next(iter(xd.values())).data_vars)
+
+    return [{"name": f"band{i + 1}"} for i in range(band_count)]
+
+
+def merge_bands(existing_bands: list[dict], new_bands: list[dict]) -> list[dict]:
+    """
+    Merge existing bands metadata with newly generated ones from xarray.
+
+    Existing bands metadata take precedence over generated ones.
+
+    :param existing_bands: existing bands metadata
+    :param new_bands: newly generated bands metadata
+    :return: merged bands metadata
+    """
+    merged = []
+
+    for i, band in enumerate(existing_bands):
+        band = dict(band)
+        band.setdefault("name", f"band{i + 1}")
+        merged.append(band)
+
+    for i in range(len(existing_bands), len(new_bands)):
+        merged.append(new_bands[i])
+
+    return merged