From 0c829a1bf3196360567da22109f040bca3511bb7 Mon Sep 17 00:00:00 2001
From: David Huard <huard.david@ouranos.ca>
Date: Fri, 29 Sep 2023 13:43:10 -0400
Subject: [PATCH] harmonized datacube and cmip extensions

---
 STACpopulator/stac_utils.py             | 26 ++++---------------------
 implementations/CMIP6-UofT/add_CMIP6.py | 16 ++++++++++++---
 2 files changed, 17 insertions(+), 25 deletions(-)

diff --git a/STACpopulator/stac_utils.py b/STACpopulator/stac_utils.py
index c26542f..24efb07 100644
--- a/STACpopulator/stac_utils.py
+++ b/STACpopulator/stac_utils.py
@@ -152,6 +152,8 @@ class Item(BaseModel):
 
 class CFJsonItem:
     """Return STAC Item from CF JSON metadata, as provided by `xncml.Dataset.to_cf_dict`."""
+    axis = {"X": "x", "Y": "y", "Z": "z", "T": "t", "longitude": "x", "latitude": "y", "vertical": "z", "time": "t"}
+
     def __init__(self, iid: str, attrs: dict, datamodel=None):
         """
         Create STAC Item from CF JSON metadata.
@@ -249,28 +251,8 @@ def ncattrs_to_bbox(self) -> list:
             float(attrs["geospatial_lat_max"][0]),
         ]
 
-
-class DatacubeExt:
-    """Extend STAC Item with Datacube properties."""
-    axis = {"X": "x", "Y": "y", "Z": "z", "T": "t", "longitude": "x", "latitude": "y", "vertical": "z", "time": "t"}
-
-    def __init__(self, obj: CFJsonItem):
-        """
-        Add Datacube extension to STAC Item.
-
-        Parameters
-        ----------
-        obj : CFJsonItem
-            STAC Item created from CF JSON metadata.
-        """
-        self.obj = obj
-        self.attrs = obj.attrs
-
-        self.ext = DatacubeExtension.ext(self.obj.item, add_if_missing=True)
-        self.ext.apply(dimensions=self.dimensions(), variables=self.variables())
-
     def dimensions(self) -> dict:
-        """Return Dimension objects."""
+        """Return Dimension objects required for Datacube extension."""
 
         dims = {}
         for name, length in self.attrs["dimensions"].items():
@@ -304,7 +286,7 @@ def dimensions(self) -> dict:
         return dims
 
     def variables(self)->dict:
-        """Return Variable objects"""
+        """Return Variable objects required for Datacube extension."""
         variables = {}
 
         for name, meta in self.attrs["variables"].items():
diff --git a/implementations/CMIP6-UofT/add_CMIP6.py b/implementations/CMIP6-UofT/add_CMIP6.py
index 1b219ee..4fdea40 100644
--- a/implementations/CMIP6-UofT/add_CMIP6.py
+++ b/implementations/CMIP6-UofT/add_CMIP6.py
@@ -6,12 +6,13 @@
 import argparse
 import pyessv
 from pydantic import AnyHttpUrl, BaseModel, Field, FieldValidationInfo, field_validator
-
+from pystac.extensions.datacube import DatacubeExtension
 
 from STACpopulator import STACpopulatorBase
+from STACpopulator.extensions import cmip6
 from STACpopulator.input import THREDDSLoader
 from STACpopulator.stac_utils import ItemProperties
-from STACpopulator.stac_utils import collection2literal, DatacubeExt, CFJsonItem
+from STACpopulator.stac_utils import collection2literal, CFJsonItem
 
 
 LOGGER = logging.getLogger(__name__)
@@ -148,8 +149,17 @@ def create_stac_item(self, item_name: str, item_data: MutableMapping[str, Any])
 
         obj = CFJsonItem(iid, item_data, self.props_model)
 
+        # Add CMIP6 extension
+        try:
+            cmip6_ext = cmip6.CMIP6Extension.ext(obj.item, add_if_missing=True)
+            cmip6_ext.apply(item_data["attributes"])
+        except:
+            LOGGER.warning(f"Failed to add CMIP6 extension to item {item_name}")
+
+        # Add datacube extension
         try:
-            DatacubeExt(obj)
+            dc_ext = DatacubeExtension.ext(obj.item, add_if_missing=True)
+            dc_ext.apply(dimensions=obj.dimensions(), variables=obj.variables())
         except:
             LOGGER.warning(f"Failed to add Datacube extension to item {item_name}")