From 47de86ff49e09e4708a69a9bbe1035dd5911e434 Mon Sep 17 00:00:00 2001
From: Tim Heap <tim.heap@csiro.au>
Date: Wed, 10 Jan 2024 14:55:37 +1100
Subject: [PATCH] Support simple SHOC datasets with no depth coordinate

This can happen after selecting a single layer, for example.

Fixes #123
---
 docs/releases/development.rst           |  3 +++
 src/emsarray/conventions/shoc.py        | 10 ++++++++--
 tests/conventions/test_cfgrid2d.py      | 11 +++++++++++
 tests/conventions/test_shoc_standard.py |  4 ++--
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/docs/releases/development.rst b/docs/releases/development.rst
index 416ffd0..7175a6f 100644
--- a/docs/releases/development.rst
+++ b/docs/releases/development.rst
@@ -11,3 +11,6 @@ Next release (in development)
 * Drop dependency on importlib_metadata.
   This was only required to support Python 3.8, which was dropped in a previous release
   (:issue:`122`, :pr:`125`).
+* Fix an error with `ShocSimple.get_all_depth_names()`
+  when the dataset had no depth coordinates
+  (:issue:`123`, :pr:`126`).
diff --git a/src/emsarray/conventions/shoc.py b/src/emsarray/conventions/shoc.py
index 8ef5640..69c9e15 100644
--- a/src/emsarray/conventions/shoc.py
+++ b/src/emsarray/conventions/shoc.py
@@ -57,7 +57,9 @@ def get_depth_name(self) -> Hashable:
         return name
 
     def get_all_depth_names(self) -> List[Hashable]:
-        return ['z_centre', 'z_grid']
+        return [
+            name for name in ['z_centre', 'z_grid']
+            if name in self.dataset.variables]
 
     def get_time_name(self) -> Hashable:
         name = 't'
@@ -123,4 +125,8 @@ def get_depth_name(self) -> Hashable:
         return name
 
     def get_all_depth_names(self) -> List[Hashable]:
-        return [self.get_depth_name()]
+        name = 'zc'
+        if name in self.dataset.variables:
+            return [name]
+        else:
+            return []
diff --git a/tests/conventions/test_cfgrid2d.py b/tests/conventions/test_cfgrid2d.py
index b73c99b..5c81c38 100644
--- a/tests/conventions/test_cfgrid2d.py
+++ b/tests/conventions/test_cfgrid2d.py
@@ -24,6 +24,7 @@
 from emsarray.conventions import get_dataset_convention
 from emsarray.conventions.grid import CFGrid2DTopology, CFGridKind
 from emsarray.conventions.shoc import ShocSimple
+from emsarray.exceptions import NoSuchCoordinateError
 from emsarray.operations import geometry
 from tests.utils import (
     DiagonalShocGrid, ShocGridGenerator, ShocLayerGenerator,
@@ -175,6 +176,16 @@ def test_varnames():
     assert dataset.ems.get_time_name() == 'time'
 
 
+def test_no_depth_coordinate():
+    dataset = make_dataset(j_size=10, i_size=10)
+    dataset = dataset.isel({'k': -1}, drop=True)
+    print(dataset)
+
+    assert dataset.ems.get_all_depth_names() == []
+    with pytest.raises(NoSuchCoordinateError):
+        dataset.ems.get_depth_name()
+
+
 @pytest.mark.parametrize(
     ['name', 'attrs'],
     [
diff --git a/tests/conventions/test_shoc_standard.py b/tests/conventions/test_shoc_standard.py
index e66dbee..ac63da7 100644
--- a/tests/conventions/test_shoc_standard.py
+++ b/tests/conventions/test_shoc_standard.py
@@ -32,7 +32,7 @@ def make_dataset(
     corner_size: int = 0,
 ) -> xarray.Dataset:
     """
-    Make a dummy SHOC simple dataset of a particular size.
+    Make a dummy SHOC standard dataset of a particular size.
     It will have a sheared grid of points located near (0, 0),
     with increasing j moving north east, and increasing i moving south east.
 
@@ -187,7 +187,7 @@ def make_dataset(
 def test_make_dataset():
     dataset = make_dataset(j_size=5, i_size=9, corner_size=2)
 
-    # Check that this is recognised as a ShocSimple dataset
+    # Check that this is recognised as a ShocStandard dataset
     assert get_dataset_convention(dataset) is ShocStandard
 
     # Check that the correct convention is used