fix issue when calling to_grib() on a dataset with single-valued dime…

…nsions. Fixes #347
ecmwf · Aug 14, 2023 · 1066dd1 · 1066dd1
1 parent 00e936e
commit 1066dd1
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -2,6 +2,13 @@
 Changelog for cfgrib
 ====================
 
+0.9.10.5 (2023-xx-xx)
+-----------------
+
+- fix issue where to_grib() could crash if given a dataset with a single-valued dimension
+  See `#347 <https://github.com/ecmwf/cfgrib/issues/347>`_.
+
+
 0.9.10.4 (2023-05-19)
 ---------------------
 

diff --git a/cfgrib/xarray_to_grib.py b/cfgrib/xarray_to_grib.py
@@ -167,7 +167,11 @@ def merge_grib_keys(grib_keys, detected_grib_keys, default_grib_keys):
 def expand_dims(data_var: xr.DataArray) -> T.Tuple[T.List[str], xr.DataArray]:
     coords_names = []  # type: T.List[str]
     for coord_name in dataset.ALL_HEADER_DIMS + ALL_TYPE_OF_LEVELS + dataset.ALL_REF_TIME_KEYS:
-        if coord_name in data_var.coords and data_var.coords[coord_name].size == 1:
+        if (
+            coord_name in data_var.coords
+            and data_var.coords[coord_name].size == 1
+            and coord_name not in data_var.dims
+        ):
             data_var = data_var.expand_dims(coord_name)
         if coord_name in data_var.dims:
             coords_names.append(coord_name)

diff --git a/tests/test_40_xarray_to_grib_regular_ll.py b/tests/test_40_xarray_to_grib_regular_ll.py
@@ -11,18 +11,18 @@
 
 from cfgrib import xarray_to_grib
 
-
-@pytest.fixture()
-def canonic_da() -> xr.DataArray:
+# we make sure to test the cases where we have a) multiple dates, and b) a single date
+@pytest.fixture(params=[4, 1])
+def canonic_da(request) -> xr.DataArray:
     coords: T.List[T.Any] = [
-        pd.date_range("2018-01-01T00:00", "2018-01-02T12:00", periods=4),
+        pd.date_range("2018-01-01T00:00", "2018-01-02T12:00", periods=request.param),
         pd.timedelta_range(0, "12h", periods=2),
         [1000.0, 850.0, 500.0],
         np.linspace(90.0, -90.0, 5),
         np.linspace(0.0, 360.0, 6, endpoint=False),
     ]
     da = xr.DataArray(
-        np.zeros((4, 2, 3, 5, 6)),
+        np.zeros((request.param, 2, 3, 5, 6)),
         coords=coords,
         dims=["time", "step", "isobaricInhPa", "latitude", "longitude"],
     )