34 bug identify icechunk store groups directly from item metadata

OceanDataStore/catalog/oceandatacatalog.py

@@ -9,8 +9,6 @@
 Authors:
     - Ollie Tooth
 """
-from typing import Optional
-
 import icechunk
 import numpy as np
 import pystac
@@ -191,9 +189,9 @@ def apply_bbox(ds: xr.Dataset,
     """
     # -- Validate Inputs -- #
     if not isinstance(ds, xr.Dataset):
-        raise ValueError("'ds' must be an xarray Dataset.")
+        raise TypeError("'ds' must be an xarray Dataset.")
     if not (isinstance(bbox, tuple) and len(bbox) == 4):
-        raise ValueError("'bbox' must be a tuple of the form (min_lon, max_lon, min_lat, max_lat).")
+        raise TypeError("'bbox' must be a tuple of the form (min_lon, max_lon, min_lat, max_lat).")
 
     # -- Identify geographical coordinate names & dimensions -- #
     # Default lat/lon coord names:
@@ -270,13 +268,11 @@ def apply_time_bounds(ds: xr.Dataset,
     """
     # -- Validate Inputs -- #
     if not isinstance(ds, xr.Dataset):
-        raise ValueError("'ds' must be an xarray Dataset.")
-    if start_datetime is not None:
-        if not isinstance(start_datetime, str):
+        raise TypeError("'ds' must be an xarray Dataset.")
+    if (start_datetime is not None) and not isinstance(start_datetime, str):
             raise ValueError("'start_datetime' must be a string in ISO format (e.g., 'YYYY-MM-DDTHH:MM:SS').")
-    if end_datetime is not None:
-        if not isinstance(end_datetime, str):
-            raise ValueError("'end_datetime' must be a string in ISO format (e.g., 'YYYY-MM-DDTHH:MM:SS').")
+    if (end_datetime is not None) and not isinstance(end_datetime, str):
+        raise ValueError("'end_datetime' must be a string in ISO format (e.g., 'YYYY-MM-DDTHH:MM:SS').")
 
     # -- Identify time dimension -- #
     for coord in ds.dims:
@@ -353,7 +349,7 @@ class OceanDataCatalog:
     """
     def __init__(self,
                  catalog_name: str = "noc-stac",
-                 catalog_url: str = None
+                 catalog_url: str | None = None
                  ):
         # Define the URL to the NOC STAC root catalog:
         self._stac_url = catalog_url or f"https://noc-msm-o.s3-ext.jc.rl.ac.uk/oceandatastore/{catalog_name}/catalog.json"
@@ -476,7 +472,7 @@ def available_items(self) -> list[str]:
         else:
             # Return all Item IDs from the current Collection or root Catalog:
             scope = self.Collection if self.Collection else self.Catalog
-            return list(item.id for item in scope.get_items(recursive=True))
+            return [item.id for item in scope.get_items(recursive=True)]
 
 
     def summary(self) -> CatalogSummary:
@@ -604,10 +600,9 @@ def item_summary(self, id: str) -> CatalogSummary:
                     item = it
                     break
         if item is None:
-            try:
-                item = self._open_item(id=id)
-            except Exception:
-                raise ValueError(f"Item '{id}' not found in Catalog.")
+            item = self._open_item(id)
+            if item is None:
+                raise RuntimeError(f"Item ID '{id}' not found in Catalog.")
 
         props    = item.properties
         title    = props.get("title", "")
@@ -791,11 +786,11 @@ def item_summary(self, id: str) -> CatalogSummary:
 
     def _filter_items(self,
                       items: list[pystac.Item],
-                      dataset_type: Optional[str] = None,
-                      product_type: Optional[str] = None,
-                      variable_name: Optional[str] = None,
-                      standard_name: Optional[str] = None,
-                      item_name: Optional[str] = None
+                      dataset_type: str | None = None,
+                      product_type: str | None = None,
+                      variable_name: str | None = None,
+                      standard_name: str | None = None,
+                      item_name: str | None = None
                       ):
         """
         Filter Items based on specified dataset type, product type,
@@ -840,12 +835,12 @@ def clear(self) -> None:
 
 
     def search(self,
-               collection: Optional[str] = None,
-               dataset_type: Optional[str] = None,
-               product_type: Optional[str] = None,
-               variable_name: Optional[str] = None,
-               standard_name: Optional[str] = None,
-               item_name: Optional[str] = None
+               collection: str | None = None,
+               dataset_type: str | None = None,
+               product_type: str | None = None,
+               variable_name: str | None = None,
+               standard_name: str | None = None,
+               item_name: str | None = None
                ) -> None:
         """
         Search the OceanDataCatalog for Items matching the specified criteria.
@@ -923,7 +918,7 @@ def search(self,
     def _open_item(
             self,
             id: str,
-        ) -> pystac.Item:
+        ) -> pystac.Item | None:
         """
         Open a STAC Item directly from the Item ID.
 
@@ -934,8 +929,8 @@ def _open_item(
 
         Returns
         -------
-        pystac.Item
-            STAC Item object.
+        pystac.Item | None
+            STAC Item object if found, otherwise None.
         """
         # Define components of Item ID path:
         parts = id.split("/")
@@ -1056,7 +1051,7 @@ def _open_zarr_store(
 
     def open_repo(self,
                   id: str,
-                  asset_key: Optional[str] = None
+                  asset_key: str | None = None
                   ) -> icechunk.Repository:
         """
         Open STAC Item asset as an Icechunk Repository.
@@ -1085,14 +1080,13 @@ def open_repo(self,
             raise TypeError("'id' must be a string.")
 
         # -- Collect Item Asset -- #
-        try:
-            item = self._open_item(id=id)
-        except Exception:
+        item = self._open_item(id)
+        if item is None:
             raise RuntimeError(f"Item ID '{id}' not found in Catalog.")
 
         # Infer asset key from Item ID if not provided:
         if asset_key is None:
-            asset_key = list(item.assets.keys())[0]
+            asset_key = next(iter(item.assets.keys()))
         asset = item.assets.get(asset_key)
         if asset is None:
             raise ValueError(f"Asset key '{asset_key}' not found in Item ID '{id}'.")
@@ -1114,14 +1108,14 @@ def open_repo(self,
 
     def open_dataset(self,
                      id: str,
-                     group: Optional[str] = None,
-                     variable_names: Optional[list[str]] = None,
-                     start_datetime: Optional[str] = None,
-                     end_datetime: Optional[str] = None,
-                     bbox: Optional[tuple[float | int, float | int, float | int, float | int]] = None,
+                     group: str | None = None,
+                     variable_names: list[str] | None = None,
+                     start_datetime: str | None = None,
+                     end_datetime: str | None = None,
+                     bbox: tuple[float | int, float | int, float | int, float | int] | None = None,
                      branch: str = "main",
                      consolidated: bool = True,
-                     asset_key: Optional[str] = None
+                     asset_key: str | None = None
                     ) -> xr.Dataset:
         """
         Open STAC Item asset as an xarray Dataset.
@@ -1132,7 +1126,7 @@ def open_dataset(self,
             Item ID to open asset.
         group : str, optional
             Group within the Zarr or Icechunk repository to read. Default is None,
-            which reads from the root of the repository.
+            which uses the group specified in the Item asset metadata.
         variable_names : list[str], optional
             List of variable names to be parsed from the dataset.
             Default is to return all variables.
@@ -1176,7 +1170,7 @@ def open_dataset(self,
             raise TypeError("'group' must be a string or None.")
         if not isinstance(variable_names, (type(None), list)):
             raise TypeError("'variable_names' must be a list of strings.")
-        if variable_names is not None and not all([isinstance(var, str) for var in variable_names]):
+        if variable_names is not None and not all(isinstance(var, str) for var in variable_names):
             raise TypeError("'variable_names' must be a list of strings.")
         if not isinstance(start_datetime, (type(None), str)):
             raise TypeError("'start_datetime' must be a string or None.")
@@ -1192,14 +1186,13 @@ def open_dataset(self,
             raise TypeError("'consolidated' must be a boolean.")
 
         # -- Collect Item Asset -- #
-        try:
-            item = self._open_item(id=id)
-        except Exception:
+        item = self._open_item(id)
+        if item is None:
             raise RuntimeError(f"Item ID '{id}' not found in Catalog.")
 
         # Infer asset key from Item ID if not provided:
         if asset_key is None:
-            asset_key = list(item.assets.keys())[0]
+            asset_key = next(iter(item.assets.keys()))
         asset = item.assets.get(asset_key)
         if asset is None:
             raise ValueError(f"Asset key '{asset_key}' not found in Item ID '{id}'.")
@@ -1208,10 +1201,13 @@ def open_dataset(self,
 
         # Open Icechunk Repository as xarray Dataset:
         if asset.to_dict()['type'] == "application/vnd.zarr+icechunk":
-            required_fields = ['bucket', 'prefix', 'anonymous', 'endpoint_url']
+            required_fields = ['bucket', 'prefix', 'group', 'anonymous', 'endpoint_url']
             for field in required_fields:
                 if field not in fields:
                     raise ValueError(f"Missing asset field '{field}' in item '{id}'.")
+            if group is None:
+                # Use default group from asset metadata when undefined:
+                group = fields['group']
             ds = self._open_icechunk_store(fields=fields, branch=branch, group=group)
 
         # Open Zarr store as xarray Dataset:

tests/unit/catalog/conftest.py

@@ -72,13 +72,13 @@ def catalog_instance(mock_catalog):
     OceanDataCatalog with pystac patched and two mock Items pre-loaded.
 
     Items:
-      - "noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/gn/T1y"  (platform gn, has tos_con/sos_con)
-      - "noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/gn/domain" (platform gn, no variables)
+      - "noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/T1y"  (platform gn, has tos_con/sos_con)
+      - "noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/domain" (platform gn, no variables)
     """
     catalog = OceanDataCatalog()
 
     item_era5 = make_mock_item(
-        item_id="noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/gn/T1y",
+        item_id="noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/T1y",
         properties={
             "platform": "gn",
             "variables": ["tos_con", "sos_con"],
@@ -94,7 +94,7 @@ def catalog_instance(mock_catalog):
         },
     )
     item_domain = make_mock_item(
-        item_id="noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/gn/domain",
+        item_id="noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/domain",
         properties={
             "platform": "gn",
             "variables": [],

tests/unit/catalog/test_oceandatacatalog.py

@@ -6,10 +6,11 @@
 Authors:
     - Ollie Tooth
 """
-import pytest
+from unittest.mock import MagicMock
+
 import numpy as np
+import pytest
 import xarray as xr
-from unittest.mock import MagicMock
 
 from OceanDataStore.catalog.oceandatacatalog import CatalogSummary
 
@@ -158,17 +159,17 @@ def test_item_summary_type_error(self, catalog_instance):
 
     def test_item_summary_found_in_items(self, catalog_instance):
         result = catalog_instance.item_summary(
-            id="noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/gn/T1y"
+            id="noc-npd-era5/npd-eorca1-era5v1/r1i1c1f1/T1y"
         )
         assert isinstance(result, CatalogSummary)
 
     def test_item_summary_invalid_id_raises(self, catalog_instance, mocker):
         catalog_instance.Items = None
         mocker.patch.object(
-            catalog_instance, "_open_item", side_effect=Exception("not found")
+            catalog_instance, "_open_item", return_value=None
         )
-        with pytest.raises(ValueError, match="Item 'nonexistent' not found in Catalog"):
-            catalog_instance.item_summary(id="nonexistent")
+        with pytest.raises(RuntimeError, match="Item ID 'invalid_id' not found in Catalog."):
+            catalog_instance.item_summary(id="invalid_id")
 
 
 class TestOceanDataCatalogOpenRepo:
@@ -178,9 +179,9 @@ def test_open_repo_invalid_id_type(self, catalog_instance):
 
     def test_open_repo_invalid_id_raises_runtime_error(self, catalog_instance, mocker):
         mocker.patch.object(
-            catalog_instance, "_open_item", side_effect=Exception("not found")
+            catalog_instance, "_open_item", return_value=None
         )
-        with pytest.raises(RuntimeError, match="Item ID 'invalid_id' not found in Catalog"):
+        with pytest.raises(RuntimeError, match="Item ID 'invalid_id' not found in Catalog."):
             catalog_instance.open_repo(id="invalid_id")
 
     def test_open_repo_invalid_asset_key(self, catalog_instance, mocker):
@@ -250,9 +251,9 @@ def test_open_dataset_invalid_asset_key(self, catalog_instance, mocker):
 
     def test_open_dataset_invalid_id_raises_runtime_error(self, catalog_instance, mocker):
         mocker.patch.object(
-            catalog_instance, "_open_item", side_effect=Exception("not found")
+            catalog_instance, "_open_item", return_value=None
         )
-        with pytest.raises(RuntimeError, match="Item ID 'invalid_id' not found in Catalog"):
+        with pytest.raises(RuntimeError, match="Item ID 'invalid_id' not found in Catalog."):
             catalog_instance.open_dataset(id="invalid_id")
 
     def test_open_dataset_invalid_variable_names(self, catalog_instance, mocker):
@@ -261,6 +262,7 @@ def test_open_dataset_invalid_variable_names(self, catalog_instance, mocker):
         mock_asset.extra_fields = {
             "bucket": "my-bucket",
             "prefix": "my-prefix",
+            "group": None,
             "anonymous": True,
             "endpoint_url": "https://example.com",
         }