Add ensure_path decorator

Author: HCookie

src/anemoi/inference/decorators.py

@@ -8,11 +8,15 @@
 # nor does it submit to any jurisdiction.
 
 
+import logging
+from pathlib import Path
 from typing import Any
 from typing import TypeVar
 
 from anemoi.inference.context import Context
 
+LOG = logging.getLogger("anemoi.inference")
+
 MARKER = object()
 F = TypeVar("F", bound=type)
 
@@ -64,3 +68,67 @@ def __init__(wrapped_cls, context: Context, main: object = MARKER, *args: Any, *
                 super().__init__(context, *args, **kwargs)
 
         return type(cls.__name__, (WrappedClass,), {})
+
+
+class ensure_path:
+    """Decorator to ensure a path argument is a Path object and optionally exists.
+
+    If `is_dir` is True, the path is treated as a directory, if not for files, the parent directory is treated as a directory.
+    If `must_exist` is True, the directory must exist.
+    If `create` is True, the directory will be created if it doesn't exist.
+
+    For example:
+    ```
+    @ensure_path("dir", create=True)
+    class GribOutput
+        def __init__(context, dir=None, archive_requests=None):
+            ...
+    """
+
+    def __init__(self, arg: str, is_dir: bool = False, create: bool = True, must_exist: bool = False):
+        self.arg = arg
+        self.is_dir = is_dir
+        self.create = create
+        self.must_exist = must_exist
+
+    def __call__(self, cls: F) -> F:
+        """Decorate the object to ensure the path argument is a Path object."""
+
+        class WrappedClass(cls):
+            def __init__(wrapped_cls, context: Context, *args: Any, **kwargs: Any) -> None:
+                if self.arg not in kwargs:
+                    LOG.debug(f"Argument '{self.arg}' not found in kwargs, cannot ensure path.")
+                    super().__init__(context, *args, **kwargs)
+                    return
+
+                path = kwargs[self.arg] = Path(kwargs[self.arg])
+                if not self.is_dir:
+                    path = path.parent
+
+                if self.must_exist:
+                    if not path.exists():
+                        raise FileNotFoundError(f"Path '{path}' does not exist.")
+                if self.create:
+                    path.mkdir(parents=True, exist_ok=True)
+
+                super().__init__(context, *args, **kwargs)
+
+        return type(cls.__name__, (WrappedClass,), {})
+
+
+class ensure_dir(ensure_path):
+    """Decorator to ensure a directory path argument is a Path object and optionally exists.
+
+    If `must_exist` is True, the directory must exist.
+    If `create` is True, the directory will be created if it doesn't exist.
+
+    For example:
+    ```
+    @ensure_dir("dir", create=True)
+    class PlotOutput
+        def __init__(context, dir=None, ...):
+            ...
+    """
+
+    def __init__(self, arg: str, create: bool = True, must_exist: bool = False):
+        super().__init__(arg, is_dir=True, create=create, must_exist=must_exist)

src/anemoi/inference/outputs/gribfile.py

@@ -23,6 +23,7 @@
 from anemoi.inference.types import FloatArray
 from anemoi.inference.types import ProcessorConfig
 
+from ..decorators import ensure_path
 from ..decorators import main_argument
 from ..grib.encoding import GribWriter
 from ..grib.encoding import check_encoding
@@ -312,6 +313,7 @@ def _patch(r: DataRequest) -> DataRequest:
 
 @output_registry.register("grib")
 @main_argument("path")
+@ensure_path("path")
 class GribFileOutput(GribIoOutput):
     """Handles grib files."""
 
@@ -367,10 +369,6 @@ def __init__(
         split_output : bool, optional
             Whether to split the output, by default True.
         """
-        path = Path(path)
-        if not path.parent.exists():
-            path.parent.mkdir(parents=True, exist_ok=True)
-
         super().__init__(
             context,
             out=path,

src/anemoi/inference/outputs/netcdf.py

@@ -18,6 +18,7 @@
 from anemoi.inference.types import ProcessorConfig
 from anemoi.inference.types import State
 
+from ..decorators import ensure_path
 from ..decorators import main_argument
 from ..output import Output
 from . import output_registry
@@ -31,6 +32,7 @@
 
 @output_registry.register("netcdf")
 @main_argument("path")
+@ensure_path("path")
 class NetCDFOutput(Output):
     """NetCDF output class."""
 
@@ -76,7 +78,7 @@ def __init__(
 
         from netCDF4 import Dataset
 
-        self.path = Path(path)
+        self.path = path
         self.ncfile: Dataset | None = None
         self.float_size = float_size
         self.missing_value = missing_value
@@ -103,7 +105,6 @@ def open(self, state: State) -> None:
             if self.ncfile is not None:
                 return
 
-        self.path.parent.mkdir(parents=True, exist_ok=True)
         # If the file exists, we may get a 'Permission denied' error
         if os.path.exists(self.path):
             os.remove(self.path)

src/anemoi/inference/outputs/plot.py

@@ -14,6 +14,7 @@
 from anemoi.utils.grib import units
 
 from anemoi.inference.context import Context
+from anemoi.inference.decorators import ensure_dir
 from anemoi.inference.decorators import main_argument
 from anemoi.inference.types import FloatArray
 from anemoi.inference.types import ProcessorConfig
@@ -44,6 +45,7 @@ def fix(lons: FloatArray) -> FloatArray:
 
 @output_registry.register("plot")
 @main_argument("dir")
+@ensure_dir("dir")
 class PlotOutput(Output):
     """Use `earthkit-plots` to plot the outputs."""
 
@@ -105,7 +107,7 @@ def __init__(
             write_initial_state=write_initial_state,
         )
 
-        self.dir = Path(dir)
+        self.dir = dir
         self.format = format
         self.variables = variables
         self.template = template
@@ -166,7 +168,6 @@ def write_step(self, state: State) -> None:
                 "variables": "_".join(self.variables or []),
             },
         )
-        self.dir.mkdir(parents=True, exist_ok=True)
         fname = self.dir / fname
 
         fig.save(fname)

src/anemoi/inference/outputs/printer.py

@@ -21,6 +21,7 @@
 from anemoi.inference.context import Context
 from anemoi.inference.types import State
 
+from ..decorators import ensure_path
 from ..decorators import main_argument
 from ..output import Output
 from . import output_registry
@@ -105,6 +106,7 @@ def print_state(
 
 @output_registry.register("printer")
 @main_argument("max_lines")
+@ensure_path("path")
 class PrinterOutput(Output):
     """Printer output class."""
 
@@ -142,9 +144,6 @@ def __init__(
         self.f = None
 
         if path is not None:
-            path = Path(path)
-            path.parent.mkdir(parents=True, exist_ok=True)
-
             self.f = open(path, "w")
             self.print = partial(print, file=self.f)
 

src/anemoi/inference/outputs/raw.py

@@ -16,6 +16,7 @@
 from anemoi.inference.types import State
 from anemoi.inference.utils.templating import render_template
 
+from ..decorators import ensure_dir
 from ..decorators import main_argument
 from ..output import Output
 from . import output_registry
@@ -25,6 +26,7 @@
 
 @output_registry.register("raw")
 @main_argument("path")
+@ensure_dir("dir")
 class RawOutput(Output):
     """Raw output class."""
 
@@ -53,7 +55,7 @@ def __init__(
             The date format string, by default "%Y%m%d%H%M%S".
         """
         super().__init__(context, variables=variables, **kwargs)
-        self.dir = Path(dir)
+        self.dir = dir
         self.template = template
         self.strftime = strftime
 
@@ -65,7 +67,7 @@ def __repr__(self) -> str:
         str
             String representation of the RawOutput object.
         """
-        return f"RawOutput({self.path})"
+        return f"RawOutput({self.dir})"
 
     def write_step(self, state: State) -> None:
         """Write the state to a compressed .npz file.
@@ -78,8 +80,6 @@ def write_step(self, state: State) -> None:
         date = state["date"]
         basetime = date - state["step"]
 
-        self.dir.mkdir(parents=True, exist_ok=True)
-
         format_info = {
             "date": date.strftime(self.strftime),
             "step": state["step"],