Update usage with new rra-tools

pyproject.toml

@@ -41,7 +41,7 @@ xarray = "^2024.11.0"
 cdsapi = "^0.7.5"
 matplotlib = "^3.8.4"
 scikit-learn = "^1.4.2"
-rra-tools = "^1.0.10"
+rra-tools = "^1.0.16"
 netcdf4 = "^1.7.2"
 pyarrow = "^16.0.0"
 gcsfs = "^2024.6.0"

src/climate_data/cli_options.py

@@ -7,9 +7,9 @@
 It also provides global variables representing the full space of valid values for these options.
 """
 
+from collections.abc import Collection
 from typing import ParamSpec, TypeVar
 
-import click
 from rra_tools.cli_tools import (
     RUN_ALL,
     ClickOption,
@@ -18,6 +18,7 @@
     with_input_directory,
     with_num_cores,
     with_output_directory,
+    with_overwrite,
     with_progress_bar,
     with_queue,
     with_verbose,
@@ -30,7 +31,7 @@
 
 
 def with_year(
-    years: list[str],
+    years: Collection[str],
     *,
     allow_all: bool = False,
 ) -> ClickOption[_P, _T]:
@@ -41,6 +42,7 @@ def with_year(
         allow_all=allow_all,
         choices=years,
         help="Year to extract data for.",
+        convert=allow_all,
     )
 
 
@@ -54,6 +56,7 @@ def with_month(
         allow_all=allow_all,
         choices=cdc.MONTHS,
         help="Month to extract data for.",
+        convert=allow_all,
     )
 
 
@@ -67,6 +70,7 @@ def with_era5_variable(
         allow_all=allow_all,
         choices=cdc.ERA5_VARIABLES,
         help="Variable to extract.",
+        convert=allow_all,
     )
 
 
@@ -80,6 +84,7 @@ def with_era5_dataset(
         allow_all=allow_all,
         choices=cdc.ERA5_DATASETS,
         help="Dataset to extract.",
+        convert=allow_all,
     )
 
 
@@ -93,6 +98,7 @@ def with_cmip6_source(
         allow_all=allow_all,
         choices=cdc.CMIP6_SOURCES,
         help="CMIP6 source to extract.",
+        convert=allow_all,
     )
 
 
@@ -106,6 +112,7 @@ def with_cmip6_experiment(
         allow_all=allow_all,
         choices=cdc.CMIP6_EXPERIMENTS,
         help="CMIP6 experiment to extract.",
+        convert=allow_all,
     )
 
 
@@ -119,11 +126,12 @@ def with_cmip6_variable(
         allow_all=allow_all,
         choices=[v.name for v in cdc.CMIP6_VARIABLES],
         help="CMIP6 variable to extract.",
+        convert=allow_all,
     )
 
 
 def with_target_variable(
-    variable_names: list[str],
+    variable_names: Collection[str],
     *,
     allow_all: bool = False,
 ) -> ClickOption[_P, _T]:
@@ -133,6 +141,7 @@ def with_target_variable(
         allow_all=allow_all,
         choices=variable_names,
         help="Variable to generate.",
+        convert=allow_all,
     )
 
 
@@ -145,6 +154,7 @@ def with_draw(
         allow_all=allow_all,
         choices=cdc.DRAWS,
         help="Draw to process.",
+        convert=allow_all,
     )
 
 
@@ -157,14 +167,7 @@ def with_scenario(
         allow_all=allow_all,
         choices=cdc.SCENARIOS,
         help="Scenario to process.",
-    )
-
-
-def with_overwrite() -> ClickOption[_P, _T]:
-    return click.option(
-        "--overwrite",
-        is_flag=True,
-        help="Overwrite existing files.",
+        convert=allow_all,
     )
 
 

src/climate_data/data.py

@@ -187,7 +187,7 @@ def load_scenario_metadata(self) -> pd.DataFrame:
 
     def save_scenario_inclusion_metadata(self, df: pd.DataFrame) -> None:
         # Need to save to our scripts directory for doc building
-        scripts_root = Path(__file__).parent.parent / "scripts"
+        scripts_root = Path(__file__).parent.parent.parent / "scripts"
         for root_dir in [self.results_metadata, scripts_root]:
             path = root_dir / "scenario_inclusion_metadata.parquet"
             if path.exists():

src/climate_data/extract/cmip6.py

@@ -32,9 +32,9 @@ def load_cmip_data(zarr_path: str) -> xr.Dataset:
 
 
 def extract_cmip6_main(
-    cmip6_variable: str,
-    cmip6_experiment: str,
     cmip6_source: str,
+    cmip6_experiment: str,
+    cmip6_variable: str,
     output_dir: str | Path,
     overwrite: bool,
 ) -> None:
@@ -95,38 +95,38 @@ def extract_cmip6_main(
 
 
 @click.command()  # type: ignore[arg-type]
-@clio.with_cmip6_variable()
-@clio.with_cmip6_experiment()
 @clio.with_cmip6_source()
+@clio.with_cmip6_experiment()
+@clio.with_cmip6_variable()
 @clio.with_output_directory(cdc.MODEL_ROOT)
 @clio.with_overwrite()
 def extract_cmip6_task(
-    cmip6_variable: str,
-    cmip6_experiment: str,
     cmip6_source: str,
+    cmip6_experiment: str,
+    cmip6_variable: str,
     output_dir: str,
     overwrite: bool,
 ) -> None:
     extract_cmip6_main(
-        cmip6_variable,
-        cmip6_experiment,
         cmip6_source,
+        cmip6_experiment,
+        cmip6_variable,
         output_dir,
         overwrite,
     )
 
 
 @click.command()  # type: ignore[arg-type]
-@clio.with_cmip6_variable(allow_all=True)
-@clio.with_cmip6_experiment(allow_all=True)
 @clio.with_cmip6_source(allow_all=True)
+@clio.with_cmip6_experiment(allow_all=True)
+@clio.with_cmip6_variable(allow_all=True)
 @clio.with_output_directory(cdc.MODEL_ROOT)
 @clio.with_queue()
 @clio.with_overwrite()
 def extract_cmip6(
-    cmip6_source: str,
-    cmip6_experiment: str,
-    cmip6_variable: str,
+    cmip6_source: list[str],
+    cmip6_experiment: list[str],
+    cmip6_variable: list[str],
     output_dir: str,
     queue: str,
     overwrite: bool,
@@ -140,27 +140,15 @@ def extract_cmip6(
     capture model inclusion criteria as it does not account for the year range avaialable
     in the data. This determiniation is made when we proccess the data in later steps.
     """
-    sources = cdc.CMIP6_SOURCES if cmip6_source == clio.RUN_ALL else [cmip6_source]
-    experiments = (
-        cdc.CMIP6_EXPERIMENTS
-        if cmip6_experiment == clio.RUN_ALL
-        else [cmip6_experiment]
-    )
-    variables = (
-        cdc.CMIP6_VARIABLES.names()
-        if cmip6_variable == clio.RUN_ALL
-        else [cmip6_variable]
-    )
-
     overwrite_arg = {"overwrite": None} if overwrite else {}
 
     jobmon.run_parallel(
         runner="cdtask",
         task_name="extract cmip6",
         node_args={
-            "cmip6-source": list(sources),
-            "cmip6-experiment": list(experiments),
-            "cmip6-variable": variables,
+            "cmip6-source": cmip6_source,
+            "cmip6-experiment": cmip6_experiment,
+            "cmip6-variable": cmip6_variable,
         },
         task_args={
             "output-dir": output_dir,