WIP: new indexes working in XarrayZarrRecipe. mypy problems
MODIFY
pangeo_forge_recipes/patterns.py
pangeo_forge_recipes/patterns.py
@@ -134,7 +134,9 @@ class FilePattern:
def __getitem__(self, indexer: FilePatternIndex) -> str:
"""Get a filename path for a particular key. """
assert len(indexer) == len(self.combine_dims)
- format_function_kwargs = {idx.name: idx.index for idx in indexer}
+ format_function_kwargs = {
+ idx.name: dim.keys[idx.index] for idx, dim in zip(indexer, self.combine_dims)
+ }
fname = self.format_function(**format_function_kwargs)
return fname

MODIFY
tests/test_patterns.py
tests/test_patterns.py
@@ -43,8 +43,10 @@ def test_pattern_from_file_sequence():
@pytest.mark.parametrize("pickle", [False, True])
def test_file_pattern_concat_merge(pickle):
- concat = ConcatDim(name="time", keys=list(range(3)))
- merge = MergeDim(name="variable", keys=["foo", "bar"])
+ times = list(range(3))
+ varnames = ["foo", "bar"]
+ concat = ConcatDim(name="time", keys=times)
+ merge = MergeDim(name="variable", keys=varnames)
def format_function(time, variable):
return f"T_{time}_V_{variable}"
@@ -65,7 +67,7 @@ def test_file_pattern_concat_merge(pickle):
assert fp.concat_sequence_lens == {"time": None}
assert len(list(fp)) == 6
for key in fp:
- fname = format_function(**{k.name: k.index for k in key})
+ expected_fname = format_function(time=times[key[1].index], variable=varnames[key[0].index])
for k in key:
if k.name == "time":
assert k.operation == CombineOp.CONCAT
@@ -73,7 +75,7 @@ def test_file_pattern_concat_merge(pickle):
if k.name == "variable":
assert k.operation == CombineOp.MERGE
assert k.sequence_len == 2
- assert fp[key] == fname
+ assert fp[key] == expected_fname
@pytest.mark.parametrize("nkeep", [1, 2])

MODIFY
tests/test_recipes.py
tests/test_recipes.py
from contextlib import nullcontext as does_not_raise
+from dataclasses import replace
from unittest.mock import patch
import pytest
@@ -7,6 +8,7 @@ import xarray as xr
# need to import this way (rather than use pytest.lazy_fixture) to make it work with dask
from pytest_lazyfixture import lazy_fixture
+from pangeo_forge_recipes.patterns import FilePattern
from pangeo_forge_recipes.recipes.base import BaseRecipe
all_recipes = [
@@ -143,9 +145,13 @@ def test_chunks(
kwargs["metadata_cache"] = None
else:
# modify file_pattern in place to remove nitems_per_file; a bit hacky
+ new_combine_dims = []
for cdim in file_pattern.combine_dims:
if hasattr(cdim, "nitems_per_file"):
- cdim.nitems_per_file = None
+ new_combine_dims.append(replace(cdim, nitems_per_file=None))
+ else:
+ new_combine_dims.append(cdim)
+ file_pattern = FilePattern(file_pattern.format_function, *new_combine_dims)
with chunk_expectation as excinfo:
rec = RecipeClass(file_pattern, **kwargs)

WIP: new indexes working in XarrayZarrRecipe. mypy problems