add copy_pruned method
MODIFY
pangeo_forge_recipes/recipes/xarray_zarr.py
pangeo_forge_recipes/recipes/xarray_zarr.py
@@ -5,7 +5,7 @@ A Pangeo Forge Recipe
import logging
import warnings
from contextlib import ExitStack, contextmanager
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
from itertools import product
from typing import Callable, Dict, List, Optional, Sequence, Tuple
@@ -14,7 +14,7 @@ import numpy as np
import xarray as xr
import zarr
-from ..patterns import FilePattern
+from ..patterns import FilePattern, prune_pattern
from ..storage import AbstractTarget, CacheFSSpecTarget, MetadataTarget, file_opener
from ..utils import (
chunk_bounds_and_conflicts,
@@ -181,6 +181,17 @@ class XarrayZarrRecipe(BaseRecipe):
if not all_chunk_keys == set([c for val in self._inputs_chunks.values() for c in val]):
raise ValueError("_inputs_chunks and _chunks_inputs don't use the same chunk keys.")
+ def copy_pruned(self, nkeep: int = 2) -> BaseRecipe:
+ """Make a copy of this recipe with a pruned file pattern.
+
+ :param nkeep: The number of items to keep from each ConcatDim sequence.
+ """
+
+ new_pattern = prune_pattern(self.file_pattern, nkeep=nkeep)
+ return replace(self, file_pattern=new_pattern)
+
+ # below here are methods that are part of recipe execution
+
def _set_target_chunks(self):
target_concat_dim_chunks = self.target_chunks.get(self._concat_dim)
if (self._nitems_per_input is None) and (target_concat_dim_chunks is None):

MODIFY
tests/test_recipes.py
tests/test_recipes.py
@@ -23,6 +23,21 @@ def test_recipe(recipe_fixture, execute_recipe):
xr.testing.assert_identical(ds_actual, ds_expected)
+@pytest.mark.parametrize("recipe_fixture", all_recipes)
+@pytest.mark.parametrize("nkeep", [1, 2])
+def test_prune_recipe(recipe_fixture, execute_recipe, nkeep):
+ """The basic recipe test. Use this as a template for other tests."""
+
+ RecipeClass, file_pattern, kwargs, ds_expected, target = recipe_fixture
+ rec = RecipeClass(file_pattern, **kwargs)
+ rec_pruned = rec.copy_pruned(nkeep=nkeep)
+ assert len(list(rec.iter_inputs())) > len(list(rec_pruned.iter_inputs()))
+ execute_recipe(rec_pruned)
+ ds_pruned = xr.open_zarr(target.get_mapper()).load()
+ nitems_per_input = list(file_pattern.nitems_per_input.values())[0]
+ assert ds_pruned.dims["time"] == nkeep * nitems_per_input
+
+
@pytest.mark.parametrize("cache_inputs", [True, False])
@pytest.mark.parametrize("copy_input_to_local_file", [True, False])
def test_recipe_caching_copying(

add copy_pruned method