From 5a9044c6086fb8bcbe720f61144789e7975a5355 Mon Sep 17 00:00:00 2001
From: EddyCMWF <Edward.Comyn-Platt@ecmwf.int>
Date: Fri, 19 Jul 2024 11:54:19 +0100
Subject: [PATCH] coords_as_attributes functionality

---
 cfgrib/dataset.py                             |   7 ++++++-
 cfgrib/xarray_plugin.py                       |   2 ++
 tests/sample-data/soil-surface-level-mix.grib | Bin 0 -> 1836 bytes
 tests/test_50_xarray_plugin.py                |  12 ++++++++++++
 4 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 tests/sample-data/soil-surface-level-mix.grib

diff --git a/cfgrib/dataset.py b/cfgrib/dataset.py
index e29136db..63009722 100644
--- a/cfgrib/dataset.py
+++ b/cfgrib/dataset.py
@@ -494,6 +494,7 @@ def build_variable_components(
     read_keys: T.Iterable[str] = (),
     time_dims: T.Sequence[str] = ("time", "step"),
     extra_coords: T.Dict[str, str] = {},
+    coords_as_attributes: T.Dict[str, str] = {},
     cache_geo_coords: bool = True,
 ) -> T.Tuple[T.Dict[str, int], Variable, T.Dict[str, Variable]]:
     data_var_attrs = enforce_unique_attributes(index, DATA_ATTRIBUTES_KEYS, filter_by_keys)
@@ -520,6 +521,9 @@ def build_variable_components(
             and "GRIB_typeOfLevel" in data_var_attrs
         ):
             coord_name = data_var_attrs["GRIB_typeOfLevel"]
+        if coord_name in coords_as_attributes and len(values) == 1:
+            data_var_attrs[f"GRIB_{coord_name}"] = values
+            continue
         coord_name_key_map[coord_name] = coord_key
         attributes = {
             "long_name": "original GRIB coordinate for key: %s(%s)" % (orig_name, coord_name),
@@ -666,6 +670,7 @@ def build_dataset_components(
     read_keys: T.Iterable[str] = (),
     time_dims: T.Sequence[str] = ("time", "step"),
     extra_coords: T.Dict[str, str] = {},
+    coords_as_attributes: T.Dict[str, str] = {},
     cache_geo_coords: bool = True,
 ) -> T.Tuple[T.Dict[str, int], T.Dict[str, Variable], T.Dict[str, T.Any], T.Dict[str, T.Any]]:
     dimensions = {}  # type: T.Dict[str, int]
@@ -692,6 +697,7 @@ def build_dataset_components(
                 read_keys=read_keys,
                 time_dims=time_dims,
                 extra_coords=extra_coords,
+                coords_as_attributes=coords_as_attributes,
                 cache_geo_coords=cache_geo_coords,
             )
         except DatasetBuildError as ex:
@@ -814,5 +820,4 @@ def open_file(
     stream = messages.FileStream(path, errors=errors)
     index_keys = compute_index_keys(time_dims, extra_coords)
     index = open_fileindex(stream, indexpath, index_keys, ignore_keys=ignore_keys, filter_by_keys=filter_by_keys)
-
     return open_from_index(index, read_keys, time_dims, extra_coords, errors=errors, **kwargs)
diff --git a/cfgrib/xarray_plugin.py b/cfgrib/xarray_plugin.py
index 4a972512..b5eea32c 100644
--- a/cfgrib/xarray_plugin.py
+++ b/cfgrib/xarray_plugin.py
@@ -105,6 +105,7 @@ def open_dataset(
         time_dims: T.Iterable[str] = ("time", "step"),
         errors: str = "warn",
         extra_coords: T.Dict[str, str] = {},
+        coords_as_attributes: T.Dict[str, str] = {},
         cache_geo_coords: bool = True,
     ) -> xr.Dataset:
         store = CfGribDataStore(
@@ -119,6 +120,7 @@ def open_dataset(
             lock=lock,
             errors=errors,
             extra_coords=extra_coords,
+            coords_as_attributes=coords_as_attributes,
             cache_geo_coords=cache_geo_coords,
         )
         with xr.core.utils.close_on_error(store):
diff --git a/tests/sample-data/soil-surface-level-mix.grib b/tests/sample-data/soil-surface-level-mix.grib
new file mode 100644
index 0000000000000000000000000000000000000000..a4dfe37b7083e8d9cc5c30ecd010fa962445314b
GIT binary patch
literal 1836
zcmZ<{@^oTg*uu!bVA7B@@qfc|AWMu9NB|KCh(ahh$tcdm!f0S%U<ejbVEE6#zy^ee
z1AzPnhB*!l>H_Q+*e^iUhj27-I}6;b6sS8<u2L{3x1{J!K|r!j#2vo?n;h>Akt)GU
z+!gE}=+p>KVPHs_6By(A!`vi^r6k8B$48{Oqxwx<PDe_XNY|C78I1}p3KM5Am;-@3
zvYY?^Z|E*yV5hR11$kQqVy=Wn1kCVW;w$0x$4AFuiD`)*gVqbv4r>QB28j@!2W&t$
zJ6D*B7$vCb__Ab7F%VJJDf&_{qhv#!fzy-P2UQL=29+fZD=^&*4Chq^>@w7LGtXav
zSe390-wDnw&Kh<O!83g3xJKBnFmo`y;8S6H#AuE39u8o{1V8bd;_6|>5^y7Pfs29Z
znvw;@J4*i4N5odt&ZthPeN*|SaROF1Z!eHZp{kpCoCW1(3ar@BppaV-oZw|)q0#D5
zdn5lwjDlN+?T;3Zq>hji78^JX7#L<v=;)|E5&ocu<50(h1#uPES6sbtBjK&YVUCv*
zo@%_Zc;xW01FM_g6r}v8wws;*2{e1;Hbnh!4se-Z`=hWRUm)R!XNT1b1BJ>DJ{I05
zOn|w8Eur&6y+f%*@`u)nc@I|INW8JhVME5IAICQw*l<K+|A8YH_S`t5faK;I@I0tZ
z)jU|i(E!fG62Pc;z(5HM3_t`F$0QpV7!<(b5Ge;Rn*nA5kO2n{a2A|_kxzkXUY*+c
zRB(@O0fPVo!?6n|T25A+SaH(f$fIM^PG2~A>HMD)ey4Mei=B`<rm!b!a@#hIV{v<|
zc4pQ4aV+ZZT4@1uXWb+i3rJ&gvj(;Esi3tT1IW#1?wwn5?$KGb^NUVgI+1aH()lBo
z=A3(Qw(a!1(;lY|>^ak;wEf7*F9*|hJZbpA@u7Rm+ANq$n}GQN4zRgdliK-I&{{eH
z<mL-p7e1Xkb2#K&-vNWYKId0mWw}ywV%zCim#3WMy0&fmf$nvimK~Qnm{2*deixUK
zd*4nmm^+)dz*s;Uo13+$olgZp76RRT=+7y)BX3q`?BH6~v|z<1wrx@;jSlCn>f4pE
nQ0ijX#VxCX=JBlV+Qd;lr(_Z5x!`+SjbJYIUjbtQX^@)%3;(<J

literal 0
HcmV?d00001

diff --git a/tests/test_50_xarray_plugin.py b/tests/test_50_xarray_plugin.py
index 16d63dd8..2009198a 100644
--- a/tests/test_50_xarray_plugin.py
+++ b/tests/test_50_xarray_plugin.py
@@ -11,6 +11,7 @@
 TEST_DATA = os.path.join(SAMPLE_DATA_FOLDER, "regular_ll_sfc.grib")
 TEST_DATA_MISSING_VALS = os.path.join(SAMPLE_DATA_FOLDER, "fields_with_missing_values.grib")
 TEST_DATA_MULTI_PARAMS = os.path.join(SAMPLE_DATA_FOLDER, "multi_param_on_multi_dims.grib")
+TEST_DATA_MULTI_LEVTYPES = os.path.join(SAMPLE_DATA_FOLDER, "soil-surface-level-mix.grib")
 
 
 def test_plugin() -> None:
@@ -164,3 +165,14 @@ def test_xr_open_dataset_file_missing_vals() -> None:
     t2 = ds["t2m"]
     assert np.isclose(np.nanmean(t2.values[0, :, :]), 268.375)
     assert np.isclose(np.nanmean(t2.values[1, :, :]), 270.716)
+
+
+def test_xr_open_dataset_coords_to_attributes() -> None:
+    ds = xr.open_dataset(
+        TEST_DATA_MULTI_LEVTYPES, engine="cfgrib", coords_as_attributes=["surface", "depthBelowLandLayer"]
+    )
+    assert "surface" not in ds.coords
+    assert "depthBelowLandLayer" not in ds.coords
+
+    assert "GRIB_surface" in ds["t2m"].attrs
+    assert "GRIB_depthBelowLandLayer" in ds["stl1"].attrs
\ No newline at end of file