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,41 @@
+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 simple class that contains all possible injection spec values
+
+    folks can import it into their recipes to pass injection spec values
+    around to custom beam routines as the example below shows
+
+    Examples:
+        ```python
+        from config import Config
+
+        @dataclass
+        class MyCustomTransform(beam.PTransform):
+
+            # this custom transform class is not listed
+            # in `pangeo_forge_recipes.injections.py:get_injection_specs()`
+            # and the instance attr will therefore not be dependency injected
+            # but we can use the `config.Config` to pass around the injections
+            target_root: None
+
+            def expand(self,
+                pcoll: beam.PCollection[Tuple[Index, xr.Dataset]],
+            ) ->  beam.PCollection[Tuple[Index, xr.Dataset]]:
+                return pcoll
+
+        config = Config()
+        recipe = (beam.Create() | MyCustomTransform(target_storage=config.target_storage))
+        ```
+    """
+
+    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",
         },