From 5441a53932af9bf94b4bedbdae85516600850e34 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Tue, 24 Oct 2023 13:10:30 +0200
Subject: [PATCH] Fix a problem with sorting datasets that have a mix of facet
 types (#2238)

---
 esmvalcore/_recipe/from_datasets.py     | 26 +++-----------
 tests/unit/recipe/test_from_datasets.py | 46 +++++++++++++++++++++----
 2 files changed, 44 insertions(+), 28 deletions(-)

diff --git a/esmvalcore/_recipe/from_datasets.py b/esmvalcore/_recipe/from_datasets.py
index 0e454c7350..8bd33fd5e9 100644
--- a/esmvalcore/_recipe/from_datasets.py
+++ b/esmvalcore/_recipe/from_datasets.py
@@ -160,23 +160,6 @@ def _group_identical_facets(variable: Mapping[str, Any]) -> Recipe:
     return result
 
 
-class _SortableDict(dict):
-    """A `dict` class that can be sorted."""
-
-    def __lt__(self, other):
-        return tuple(self.items()) < tuple(other.items())
-
-
-def _change_dict_type(item, dict_type):
-    """Change the dict type in a nested structure."""
-    change_dict_type = partial(_change_dict_type, dict_type=dict_type)
-    if isinstance(item, dict):
-        return dict_type((k, change_dict_type(v)) for k, v in item.items())
-    if isinstance(item, (list, tuple, set)):
-        return type(item)(change_dict_type(elem) for elem in item)
-    return item
-
-
 def _group_ensemble_members(dataset_facets: Iterable[Facets]) -> list[Facets]:
     """Group ensemble members.
 
@@ -185,13 +168,15 @@ def _group_ensemble_members(dataset_facets: Iterable[Facets]) -> list[Facets]:
     """
 
     def grouper(facets):
-        return tuple((k, facets[k]) for k in sorted(facets) if k != 'ensemble')
+        return sorted(
+            (f, str(v)) for f, v in facets.items() if f != 'ensemble')
 
     result = []
-    dataset_facets = _change_dict_type(dataset_facets, _SortableDict)
     dataset_facets = sorted(dataset_facets, key=grouper)
-    for group_facets, group in itertools.groupby(dataset_facets, key=grouper):
+    for _, group_iter in itertools.groupby(dataset_facets, key=grouper):
+        group = list(group_iter)
         ensembles = [f['ensemble'] for f in group if 'ensemble' in f]
+        group_facets = group[0]
         if not ensembles:
             result.append(dict(group_facets))
         else:
@@ -199,7 +184,6 @@ def grouper(facets):
                 facets = dict(group_facets)
                 facets['ensemble'] = ensemble
                 result.append(facets)
-    result = _change_dict_type(result, dict)
     return result
 
 
diff --git a/tests/unit/recipe/test_from_datasets.py b/tests/unit/recipe/test_from_datasets.py
index c2032724a4..aacca5117d 100644
--- a/tests/unit/recipe/test_from_datasets.py
+++ b/tests/unit/recipe/test_from_datasets.py
@@ -8,7 +8,6 @@
     _group_ensemble_names,
     _group_identical_facets,
     _move_one_level_up,
-    _SortableDict,
     _to_frozen,
     datasets_to_recipe,
 )
@@ -137,12 +136,6 @@ def test_supplementary_datasets_to_recipe():
     assert datasets_to_recipe([dataset]) == recipe
 
 
-def test_sortable_dict():
-    assert _SortableDict({'a': 1}) < _SortableDict({'a': 2})
-    assert _SortableDict({'a': 1}) < _SortableDict({'a': 1, 'b': 1})
-    assert _SortableDict({'a': 1}) < _SortableDict({'b': 1})
-
-
 def test_datasets_to_recipe_group_ensembles():
     datasets = [
         Dataset(
@@ -253,6 +246,45 @@ def test_group_ensemble_members():
     ]
 
 
+def test_group_ensemble_members_mix_of_versions():
+    datasets = [
+        Dataset(
+            dataset='dataset1',
+            ensemble='r1i1p1f1',
+            exp=['historical', 'ssp585'],
+            version='v1',
+        ),
+        Dataset(
+            dataset='dataset1',
+            ensemble='r2i1p1f1',
+            exp=['historical', 'ssp585'],
+            version='v1',
+        ),
+        Dataset(
+            dataset='dataset1',
+            ensemble='r3i1p1f1',
+            exp=['historical', 'ssp585'],
+            version=['v1', 'v2'],
+        ),
+    ]
+    result = _group_ensemble_members(ds.facets for ds in datasets)
+    print(result)
+    assert result == [
+        {
+            'dataset': 'dataset1',
+            'ensemble': 'r3i1p1f1',
+            'exp': ['historical', 'ssp585'],
+            'version': ['v1', 'v2'],
+        },
+        {
+            'dataset': 'dataset1',
+            'ensemble': 'r(1:2)i1p1f1',
+            'exp': ['historical', 'ssp585'],
+            'version': 'v1',
+        },
+    ]
+
+
 def test_group_ensembles_cmip5():
     ensembles = [
         "r1i1p1",