use decorator for method wrapping
MODIFY
pangeo_forge/recipes/base.py
pangeo_forge/recipes/base.py
from abc import ABC, abstractmethod
+from functools import partial
from typing import Callable, Hashable, Iterable
from rechunker.types import MultiStagePipeline, ParallelPipelines, Stage
@@ -93,3 +94,20 @@ class BaseRecipe(ABC):
# just intercept the __post_init__ calls so they
# aren't relayed to `object`
pass
+
+
+def closure(func: Callable) -> Callable:
+ """Wrap a method to eliminate the self keyword from its signature."""
+
+ # tried using @functools.wraps, but could not get it to work right
+ def wrapped(*args, **kwargs):
+ self = args[0]
+ if len(args) > 1:
+ args = args[1:]
+ else:
+ args = ()
+ new_func = partial(func, self, *args, **kwargs)
+ new_func.__name__ = getattr(func, "__name__", None)
+ return new_func
+
+ return wrapped

MODIFY
pangeo_forge/recipes/xarray_zarr.py
pangeo_forge/recipes/xarray_zarr.py
@@ -26,7 +26,7 @@ from ..utils import (
fix_scalar_attr_encoding,
lock_for_conflicts,
)
-from .base import BaseRecipe
+from .base import BaseRecipe, closure
# use this filename to store global recipe metadata in the metadata_cache
# it will be written once (by prepare_target) and read many times (by store_chunk)
@@ -244,10 +244,9 @@ class XarrayZarrRecipe(BaseRecipe):
else:
self._concat_dim_chunks = self._nitems_per_input * self.inputs_per_chunk
- @property
- def prepare_target(self) -> Callable:
- def _prepare_target():
-
+ @property # type: ignore
+ @closure
+ def prepare_target(self) -> None:
try:
ds = self.open_target()
logger.info("Found an existing dataset in target")
@@ -275,16 +274,12 @@ class XarrayZarrRecipe(BaseRecipe):
for dim, chunk in self.target_chunks.items()
if dim in this_var.dims
}
-
- chunks = tuple(
+ encoding_chunks = tuple(
chunks.get(n, s) for n, s in enumerate(this_var.shape)
)
- encoding_chunks = chunks
else:
encoding_chunks = ds[v].shape
- logger.debug(
- f"Setting variable {v} encoding chunks to {encoding_chunks}"
- )
+ logger.debug(f"Setting variable {v} encoding chunks to {encoding_chunks}")
ds[v].encoding["chunks"] = encoding_chunks
# load all variables that don't have the sequence dim in them
@@ -296,7 +291,7 @@ class XarrayZarrRecipe(BaseRecipe):
ds[v].load()
target_mapper = self.target.get_mapper()
- logger.info(f"Storing dataset in {self.target.root_path}")
+ logger.info(f"Storing dataset in {self.target.root_path}") # type: ignore
logger.debug(f"{ds}")
with warnings.catch_warnings():
warnings.simplefilter(
@@ -319,11 +314,10 @@ class XarrayZarrRecipe(BaseRecipe):
logger.info("Caching global metadata")
meta_mapper[_GLOBAL_METADATA_KEY] = json.dumps(recipe_meta).encode("utf-8")
- return _prepare_target
-
- @property
- def cache_input(self) -> Callable:
- def cache_func(input_key: InputKey) -> None:
+ # TODO: figure out how to make mypy happy with this convoluted structure
+ @property # type: ignore
+ @closure
+ def cache_input(self, input_key: InputKey) -> None: # type: ignore
logger.info(f"Caching input {input_key}")
fname = self.file_pattern[input_key]
@@ -342,16 +336,12 @@ class XarrayZarrRecipe(BaseRecipe):
if self._cache_metadata:
self.cache_input_metadata(input_key)
- return cache_func
-
- @property
- def store_chunk(self) -> Callable:
- def _store_chunk(chunk_key: ChunkKey):
+ @property # type: ignore
+ @closure
+ def store_chunk(self, chunk_key: ChunkKey) -> None: # type: ignore
with self.open_chunk(chunk_key) as ds_chunk:
# writing a region means that all the variables MUST have concat_dim
- to_drop = [
- v for v in ds_chunk.variables if self._concat_dim not in ds_chunk[v].dims
- ]
+ to_drop = [v for v in ds_chunk.variables if self._concat_dim not in ds_chunk[v].dims]
ds_chunk = ds_chunk.drop_vars(to_drop)
target_mapper = self.target.get_mapper()
@@ -384,18 +374,14 @@ class XarrayZarrRecipe(BaseRecipe):
)
zarr_array[zarr_region] = data
- return _store_chunk
-
- @property
- def finalize_target(self) -> Callable:
- def _finalize_target():
+ @property # type: ignore
+ @closure
+ def finalize_target(self) -> None:
if self.consolidate_zarr:
logger.info("Consolidating Zarr metadata")
target_mapper = self.target.get_mapper()
zarr.consolidate_metadata(target_mapper)
- return _finalize_target
-
@contextmanager
def input_opener(self, fname: str):
try:

use decorator for method wrapping