Skip to content

Commit

Permalink
Fix a problem with sorting datasets that have a mix of facet types (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
bouweandela authored Oct 24, 2023
1 parent e3cc51b commit 5441a53
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 28 deletions.
26 changes: 5 additions & 21 deletions esmvalcore/_recipe/from_datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -185,21 +168,22 @@ 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:
for ensemble in _group_ensemble_names(ensembles):
facets = dict(group_facets)
facets['ensemble'] = ensemble
result.append(facets)
result = _change_dict_type(result, dict)
return result


Expand Down
46 changes: 39 additions & 7 deletions tests/unit/recipe/test_from_datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
_group_ensemble_names,
_group_identical_facets,
_move_one_level_up,
_SortableDict,
_to_frozen,
datasets_to_recipe,
)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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",
Expand Down

0 comments on commit 5441a53

Please sign in to comment.