clean

Author: floriankrb

src/anemoi/datasets/create/__init__.py

@@ -256,7 +256,8 @@ def check_name(
         resolution: str,
         dates: list[datetime.datetime],
         frequency: datetime.timedelta,
-        raise_exception: bool = False,
+        raise_exception: bool = True,
+        is_test: bool = False,
     ) -> None:
         """Check the name of the dataset.
 
@@ -270,13 +271,15 @@ def check_name(
             The frequency of the dataset.
         raise_exception : bool, optional
             Whether to raise an exception if the name is invalid.
+        is_test : bool, optional
+            Whether this is a test.
         """
         basename, _ = os.path.splitext(os.path.basename(self.path))
         try:
             DatasetName(basename, resolution, dates[0], dates[-1], frequency).raise_if_not_valid()
         except Exception as e:
-            if raise_exception:
-                raise
+            if raise_exception and not is_test:
+                raise e
             else:
                 LOG.warning(f"Dataset name error: {e}")
 
@@ -574,6 +577,7 @@ def __init__(
         use_threads: bool = False,
         statistics_temp_dir: str | None = None,
         progress: Any = None,
+        test: bool = False,
         cache: str | None = None,
         **kwargs: Any,
     ):
@@ -595,6 +599,8 @@ def __init__(
             The directory for temporary statistics.
         progress : Any, optional
             The progress indicator.
+        test : bool, optional
+            Whether this is a test.
         cache : Optional[str], optional
             The cache directory.
         """
@@ -607,8 +613,9 @@ def __init__(
         self.use_threads = use_threads
         self.statistics_temp_dir = statistics_temp_dir
         self.progress = progress
+        self.test = test
 
-        self.main_config = loader_config(config)
+        self.main_config = loader_config(config, is_test=test)
 
         # self.registry.delete() ??
         self.tmp_statistics.delete()
@@ -741,6 +748,7 @@ def _run(self) -> int:
 
         self.dataset.check_name(
             raise_exception=self.check_name,
+            is_test=self.test,
             resolution=resolution,
             dates=dates,
             frequency=frequency,

src/anemoi/datasets/create/config.py

@@ -18,6 +18,8 @@
 from anemoi.utils.config import load_any_dict_format
 from earthkit.data.core.order import normalize_order_by
 
+from anemoi.datasets.dates.groups import Groups
+
 LOG = logging.getLogger(__name__)
 
 
@@ -338,20 +340,72 @@ def _prepare_serialisation(o: Any) -> Any:
     return str(o)
 
 
-def loader_config(config: dict) -> LoadersConfig:
+def set_to_test_mode(cfg: dict) -> None:
+    """Modifies the configuration to run in test mode.
+
+    Parameters
+    ----------
+    cfg : dict
+        The configuration dictionary.
+    """
+    NUMBER_OF_DATES = 4
+
+    LOG.warning(f"Running in test mode. Changing the list of dates to use only {NUMBER_OF_DATES}.")
+    groups = Groups(**LoadersConfig(cfg).dates)
+
+    dates = groups.provider.values
+    cfg["dates"] = dict(
+        start=dates[0],
+        end=dates[NUMBER_OF_DATES - 1],
+        frequency=groups.provider.frequency,
+        group_by=NUMBER_OF_DATES,
+    )
+
+    num_ensembles = count_ensembles(cfg)
+
+    def set_element_to_test(obj):
+        if isinstance(obj, (list, tuple)):
+            for v in obj:
+                set_element_to_test(v)
+            return
+        if isinstance(obj, (dict, DotDict)):
+            if "grid" in obj and num_ensembles > 1:
+                previous = obj["grid"]
+                obj["grid"] = "20./20."
+                LOG.warning(f"Running in test mode. Setting grid to {obj['grid']} instead of {previous}")
+            if "number" in obj and num_ensembles > 1:
+                if isinstance(obj["number"], (list, tuple)):
+                    previous = obj["number"]
+                    obj["number"] = previous[0:3]
+                    LOG.warning(f"Running in test mode. Setting number to {obj['number']} instead of {previous}")
+            for k, v in obj.items():
+                set_element_to_test(v)
+            if "constants" in obj:
+                constants = obj["constants"]
+                if "param" in constants and isinstance(constants["param"], list):
+                    constants["param"] = ["cos_latitude"]
+
+    set_element_to_test(cfg)
+
+
+def loader_config(config: dict, is_test: bool = False) -> LoadersConfig:
     """Loads and validates the configuration for dataset loaders.
 
     Parameters
     ----------
     config : dict
         The configuration dictionary.
+    is_test : bool, optional
+        Whether to run in test mode. Defaults to False.
 
     Returns
     -------
     LoadersConfig
         The validated configuration object.
     """
     config = Config(config)
+    if is_test:
+        set_to_test_mode(config)
     obj = LoadersConfig(config)
 
     # yaml round trip to check that serialisation works as expected
@@ -372,9 +426,6 @@ def loader_config(config: dict) -> LoadersConfig:
             LOG.info(f"Setting env variable {k}={v}")
             os.environ[k] = str(v)
 
-    # Used by pytest only
-    # copy.pop('checks', None)
-
     return copy