diff --git a/src/anemoi/registry/commands/datasets.py b/src/anemoi/registry/commands/datasets.py index 72bfaaa..49966af 100644 --- a/src/anemoi/registry/commands/datasets.py +++ b/src/anemoi/registry/commands/datasets.py @@ -39,6 +39,7 @@ def add_arguments(self, command_parser): command_parser.add_argument("--set-status", help="Set the status to the dataset") command_parser.add_argument("--add-location", nargs="+", help="Add a location to the dataset") + command_parser.add_argument("--add-recipe", help="Add a recipe file") def check_arguments(self, args): pass @@ -47,6 +48,7 @@ def run_from_identifier( self, identifier, add_location, + add_recipe, set_status, unregister, json, @@ -65,6 +67,8 @@ def run_from_identifier( entry.remove_location(**remove_location) if set_status: entry.set_status(set_status) + if add_recipe: + entry.add_recipe(add_recipe) if json: print(entry.as_json()) @@ -75,6 +79,7 @@ def run_from_path( register, unregister, add_location, + add_recipe, json, set_status, # remove_location, @@ -98,6 +103,8 @@ def run_from_path( entry.set_status(set_status) # if delete: # entry.delete() + if add_recipe: + entry.add_recipe(add_recipe) if json: print(entry.as_json()) diff --git a/src/anemoi/registry/entry/dataset.py b/src/anemoi/registry/entry/dataset.py index 7c1155b..bd2692f 100644 --- a/src/anemoi/registry/entry/dataset.py +++ b/src/anemoi/registry/entry/dataset.py @@ -8,6 +8,7 @@ import logging import os +import yaml from anemoi.datasets import open_dataset from . import CatalogueEntry @@ -27,6 +28,16 @@ def add_location(self, platform, path): patch = [{"op": "add", "path": f"/locations/{platform}", "value": {"path": path}}] self.patch(patch) + def add_recipe(self, file): + if not os.path.exists(file): + raise FileNotFoundError(f"Recipe file not found: {file}") + if not file.endswith(".yaml"): + LOG.warning("Recipe file extension is not .yaml") + with open(file) as f: + recipe = yaml.safe_load(f) + patch = [{"op": "add", "path": "/recipe", "value": recipe}] + self.patch(patch) + def load_from_path(self, path): import zarr