Add Config to Pass Around Injected Args

examples/feedstock/gpcp_from_gcs_w_custom_transform.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass, field
+from typing import Union
+
+import apache_beam as beam
+import pandas as pd
+import zarr
+
+from pangeo_forge_recipes.config import Config
+from pangeo_forge_recipes.patterns import ConcatDim, FilePattern
+from pangeo_forge_recipes.storage import FSSpecTarget
+from pangeo_forge_recipes.transforms import (
+    OpenURLWithFSSpec,
+    OpenWithXarray,
+    RequiredAtRuntimeDefault,
+    StoreToZarr,
+)
+
+dates = [
+    d.to_pydatetime().strftime("%Y%m%d")
+    for d in pd.date_range("1996-10-01", "1999-02-01", freq="D")
+]
+
+
+def make_url(time):
+    url_base = "https://storage.googleapis.com/pforge-test-data"
+    return f"{url_base}/gpcp/v01r03_daily_d{time}.nc"
+
+
+concat_dim = ConcatDim("time", dates, nitems_per_file=1)
+pattern = FilePattern(make_url, concat_dim)
+
+
+def test_ds(store: zarr.storage.FSStore) -> zarr.storage.FSStore:
+    # This fails integration test if not imported here
+    # TODO: see if --setup-file option for runner fixes this
+    import xarray as xr
+
+    ds = xr.open_dataset(store, engine="zarr", chunks={})
+    assert ds.title == (
+        "Global Precipitation Climatatology Project (GPCP) " "Climate Data Record (CDR), Daily V1.3"
+    )
+    # Making sure that the native chunking is different from the dynamic chunking
+    assert ds.chunks["time"][0] == 1
+
+    return store
+
+
+@dataclass
+class MyCustomTransform(beam.PTransform):
+
+    target_root: Union[str, FSSpecTarget, RequiredAtRuntimeDefault] = field(
+        default_factory=RequiredAtRuntimeDefault
+    )
+
+    def expand(
+        self,
+        pcoll: beam.PCollection,
+    ) -> beam.PCollection:
+        return pcoll
+
+
+config = Config()
+
+recipe = (
+    beam.Create(pattern.items())
+    | OpenURLWithFSSpec()
+    | OpenWithXarray(file_type=pattern.file_type, xarray_open_kwargs={"decode_coords": "all"})
+    | MyCustomTransform(target_root=config.target_storage)
+    | StoreToZarr(
+        store_name="gpcp.zarr",
+        combine_dims=pattern.combine_dim_keys,
+    )
+    | "Test dataset" >> beam.Map(test_ds)
+)

examples/feedstock/meta.yaml

@@ -1,6 +1,8 @@
 recipes:
   - id: "gpcp-from-gcs"
     object: "gpcp_from_gcs:recipe"
+  - id: "gpcp-from-gcs-w-custom-transform"
+    object: "gpcp_from_gcs_w_custom_transform:recipe"
   - id: "gpcp-from-gcs-dynamic-chunks"
     object: "gpcp_from_gcs_dynamic_chunks:recipe"
   - id: "noaa-oisst"

pangeo_forge_recipes/config.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+from pangeo_forge_recipes.storage import CacheFSSpecTarget, FSSpecTarget
+from pangeo_forge_recipes.transforms import RequiredAtRuntimeDefault
+
+
+@dataclass
+class Config:
+    """a class that contains all dependency injections needed for transforms"""
+
+    target_storage: Union[str, FSSpecTarget, RequiredAtRuntimeDefault] = field(
+        default_factory=RequiredAtRuntimeDefault
+    )
+    input_cache_storage: Optional[Union[str, CacheFSSpecTarget]] = ""

pangeo_forge_recipes/injections.py

@@ -1,5 +1,9 @@
 def get_injection_specs():
     return {
+        "Config": {
+            "target_storage": "TARGET_STORAGE",
+            "input_cache_storage": "INPUT_CACHE_STORAGE",
+        },
         "StoreToZarr": {
             "target_root": "TARGET_STORAGE",
         },