Modern path handling

COSIPY.py

@@ -22,11 +22,14 @@
 
 Correspondence: [email protected]
 """
-import cProfile
+
+from __future__ import annotations
+
 import logging
-import os
 from datetime import datetime
 from itertools import product
+from pathlib import Path
+from typing import Union
 
 import numpy as np
 import pandas as pd
@@ -116,8 +119,8 @@ def main():
         #encoding[var] = dict(zlib=True, complevel=compression_level, dtype=dtype, scale_factor=scale_factor, add_offset=add_offset, _FillValue=FillValue)
         encoding[var] = dict(zlib=True, complevel=Config.compression_level)
     output_netcdf = set_output_netcdf_path()
-    output_path = create_data_directory(path='output')
-    IO.get_result().to_netcdf(os.path.join(output_path,output_netcdf), encoding=encoding, mode='w')
+    output_path = create_data_directory(name="output")
+    IO.get_result().to_netcdf(output_path / output_netcdf, encoding=encoding, mode="w")
 
     encoding = dict()
     for var in IO.get_restart().data_vars:
@@ -129,8 +132,8 @@ def main():
         #encoding[var] = dict(zlib=True, complevel=compression_level, dtype=dtype, scale_factor=scale_factor, add_offset=add_offset, _FillValue=FillValue)
         encoding[var] = dict(zlib=True, complevel=Config.compression_level)
 
-    restart_path = create_data_directory(path='restart')
-    IO.get_restart().to_netcdf(os.path.join(restart_path,f'restart_{timestamp}.nc'), encoding=encoding)
+    restart_path = create_data_directory(name='restart')
+    IO.get_restart().to_netcdf(restart_path / f'restart_{timestamp}.nc', encoding=encoding)
 
     #-----------------------------------------------
     # Stop time measurement
@@ -326,19 +329,29 @@ def run_cosipy(cluster, IO, DATA, RESULT, RESTART, futures):
 
         if Config.stake_evaluation:
             # Save the statistics and the mass balance simulations at the stakes to files
-            output_path = create_data_directory(path='output')
-            df_stat.to_csv(os.path.join(output_path,'stake_statistics.csv'),sep='\t', float_format='%.2f')
-            df_val.to_csv(os.path.join(output_path,'stake_simulations.csv'),sep='\t', float_format='%.2f')
-
-
-def create_data_directory(path: str) -> str:
+            output_path = create_data_directory(name="output")
+            df_stat.to_csv(
+                output_path / "stake_statistics.csv",
+                sep="\t",
+                float_format="%.2f",
+            )
+            df_val.to_csv(
+                output_path / "stake_simulations.csv",
+                sep="\t",
+                float_format="%.2f",
+            )
+
+
+def create_data_directory(name: Union[Path, str]) -> Path:
     """Create a directory in the configured data folder.
 
     Returns:
         Path to the created directory.
     """
-    dir_path = os.path.join(Config.data_path, path)
-    os.makedirs(dir_path, exist_ok=True)
+    if isinstance(name, Path):
+        name = name.name
+    dir_path = Path(Config.data_path) / str(name)
+    dir_path.mkdir(parents=True, exist_ok=True)
 
     return dir_path
 
@@ -355,25 +368,23 @@ def get_timestamp_label(timestamp: str) -> str:
     return (timestamp[0:10]).replace("-", "")
 
 
-def set_output_netcdf_path() -> str:
+def set_output_netcdf_path() -> Path:
     """Set the file path for the output netCDF file.
 
     Returns:
         The path to the output netCDF file.
     """
     time_start = get_timestamp_label(timestamp=Config.time_start)
     time_end = get_timestamp_label(timestamp=Config.time_end)
-    output_path = f"{Config.output_prefix}_{time_start}-{time_end}.nc"
-
-    return output_path
+    return Path(f"{Config.output_prefix}_{time_start}-{time_end}.nc")
 
 
 def start_logging():
-    """Start the python logging"""
-
-    if os.path.exists('./cosipy.yaml'):
-        with open('./cosipy.yaml', 'rt') as f:
-            config = yaml.load(f.read(),Loader=yaml.SafeLoader)
+    """Start the python logging."""
+    log_config_path = Path("./cosipy.yaml")
+    if log_config_path.exists():
+        with log_config_path.open() as f:
+            config = yaml.load(f.read(), Loader=yaml.SafeLoader)
         logging.config.dictConfig(config)
     else:
         logging.basicConfig(level=logging.INFO)

convert_config.py

@@ -19,6 +19,7 @@
 import configparser
 import inspect
 import sys
+from pathlib import Path
 
 import config
 import constants
@@ -284,12 +285,12 @@ def get_utilities_params() -> dict:
     return params
 
 
-def write_toml(parameters: dict, filename: str):
+def write_toml(parameters: dict, filename: Path | str):
     """Write parameters to .toml file."""
-
-    with open(f"{filename}.toml", "w") as f:
+    if isinstance(filename, str):
+        filename = Path(filename)
+    with filename.with_suffix(".toml").open("w") as f:
         toml.dump(parameters, f)
-
 
     print(f"Generated {filename}.toml")
 
@@ -308,15 +309,22 @@ def main():
 
     print_warning()
 
-    script_path = inspect.getfile(inspect.currentframe())
-    toml_suffix = script_path.split("/")[-2]  # avoid overwrite
+    frame = inspect.currentframe()
+    if frame is None:
+        msg = "Could not find the current frame. This is likely due to a bug in the code."
+        raise RuntimeError(msg)
+    try:
+        script_path = Path(inspect.getfile(frame)).resolve()
+    finally:
+        del frame
+    _ = script_path.parent.name  # HACK: avoid overwrite (Why is this here?)
 
     config_params = get_config_params()
-    write_toml(parameters=config_params, filename=f"config")
+    write_toml(parameters=config_params, filename="config")
     constants_params = get_constants_params()
-    write_toml(parameters=constants_params, filename=f"constants")
+    write_toml(parameters=constants_params, filename="constants")
     utilities_params = get_utilities_params()
-    write_toml(parameters=utilities_params, filename=f"utilities_config")
+    write_toml(parameters=utilities_params, filename="utilities_config")
 
 
 if __name__ == "__main__":

cosipy/config.py

@@ -1,20 +1,18 @@
-"""
-Hook configuration files for COSIPY.
-"""
+"""Hook configuration files for COSIPY."""
 
 import argparse
 import os
-import pathlib
 import sys
 from importlib.metadata import entry_points
+from pathlib import Path
 
 if sys.version_info >= (3, 11):
     import tomllib
 else:
     import tomli as tomllib  # backwards compatibility
 
 
-def get_cosipy_path_from_env(name: str = "COSIPY_DIR") -> pathlib.Path:
+def get_cosipy_path_from_env(name: str = "COSIPY_DIR") -> Path:
     """Get path to COSIPY directory.
 
     When using WRFxCSPY, the coupler will default to searching for
@@ -32,26 +30,32 @@ def get_cosipy_path_from_env(name: str = "COSIPY_DIR") -> pathlib.Path:
     Raises:
         NotADirectoryError: Invalid path.
     """
-    cosipy_path = pathlib.Path(os.environ.get(name, os.getcwd()))
+    cosipy_path = Path(os.environ.get(name, Path.cwd()))
     if not cosipy_path.is_dir():
         raise NotADirectoryError(f"Invalid path at: {cosipy_path}")
 
     return cosipy_path
 
 
+cosipy_path = get_cosipy_path_from_env()
+default_config_path = cosipy_path / "config.toml"
+default_slurm_path = cosipy_path / "slurm_config.toml"
+default_constants_path = cosipy_path / "constants.toml"
+default_utilities_path = cosipy_path / "utilities_config.toml"
+
+
 def set_parser() -> argparse.ArgumentParser:
     """Set argument parser for COSIPY."""
     tagline = "Coupled snowpack and ice surface energy and mass balance model in Python."
     parser = argparse.ArgumentParser(prog="COSIPY", description=tagline)
-    cosipy_path = get_cosipy_path_from_env()
 
     # Optional arguments
     parser.add_argument(
         "-c",
         "--config",
-        default=cosipy_path / "config.toml",
+        default=default_config_path,
         dest="config_path",
-        type=pathlib.Path,
+        type=Path,
         metavar="<path>",
         required=False,
         help="relative path to configuration file",
@@ -60,9 +64,9 @@ def set_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "-x",
         "--constants",
-        default=cosipy_path / "constants.toml",
+        default=default_constants_path,
         dest="constants_path",
-        type=pathlib.Path,
+        type=Path,
         metavar="<path>",
         required=False,
         help="relative path to constants file",
@@ -71,9 +75,9 @@ def set_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "-s",
         "--slurm",
-        default=cosipy_path / "slurm_config.toml",
+        default=default_slurm_path,
         dest="slurm_path",
-        type=pathlib.Path,
+        type=Path,
         metavar="<path>",
         required=False,
         help="relative path to Slurm configuration file",
@@ -153,7 +157,7 @@ class TomlLoader(object):
     """Load and parse configuration files."""
 
     @staticmethod
-    def get_raw_toml(file_path: str = "./config.toml") -> dict:
+    def get_raw_toml(file_path: Path = default_config_path) -> dict:
         """Open and load .toml configuration file.
 
         Args:
@@ -162,7 +166,7 @@ def get_raw_toml(file_path: str = "./config.toml") -> dict:
         Returns:
             Loaded .toml data.
         """
-        with open(file_path, "rb") as f:
+        with file_path.open("rb") as f:
             raw_config = tomllib.load(f)
 
         return raw_config
@@ -191,7 +195,7 @@ def __init__(self):
         self.load(self.args.config_path)
 
     @classmethod
-    def load(cls, path: str = "./config.toml"):
+    def load(cls, path: Path = default_config_path):
         raw_toml = cls.get_raw_toml(path)
         parsed_toml = cls.set_correct_config(raw_toml)
         cls.set_config_values(parsed_toml)
@@ -245,7 +249,7 @@ def __init__(self):
         self.load(self.args.slurm_path)
 
     @classmethod
-    def load(cls, path: str = "./slurm_config.toml"):
+    def load(cls, path: Path = default_slurm_path):
         raw_toml = cls.get_raw_toml(path)
         parsed_toml = cls.set_correct_config(raw_toml)
         cls.set_config_values(parsed_toml)

cosipy/constants.py

@@ -1,4 +1,6 @@
 import sys
+from pathlib import Path
+from typing import Literal
 
 from cosipy.config import Config, TomlLoader, get_user_arguments
 

cosipy/cpkernel/io.py

@@ -1,10 +1,11 @@
-"""
- Read the input data (model forcing) and write the output to netCDF file.
-"""
+"""Read the input data (model forcing) and write the output to netCDF file."""
+
+from __future__ import annotations
 
 import os
 import warnings
 from datetime import datetime
+from pathlib import Path
 
 import numpy as np
 import xarray as xr
@@ -14,8 +15,7 @@
 
 
 class IOClass:
-
-    def __init__(self, DATA=None):
+    def __init__(self, DATA: xr.Dataset | None = None):
         """Initialise the IO Class.
 
         Attributes:
@@ -138,11 +138,15 @@ def create_data_file(self) -> xr.Dataset:
             start_timestamp = self.get_datetime(time_start)
             end_timestamp = self.get_datetime(time_end)
             timestamp = start_timestamp.strftime("%Y-%m-%dT%H-%M")
-            restart_path = os.path.join(
-                Config.data_path, "restart", f"restart_{timestamp}.nc"
-            )
+            restart_path = Path(Config.data_path) / "restart" / f"restart_{timestamp}.nc"
+            if not restart_path.is_file():
+                msg = f"No restart file available at {restart_path}"
+                raise FileNotFoundError(msg)
+            if start_timestamp == end_timestamp:
+                msg = f"Start date {time_start} equals end date {time_end}"
+                raise IndexError(msg)
             try:
-                if not os.path.isfile(restart_path):
+                if not restart_path.is_file():
                     raise FileNotFoundError
                 elif start_timestamp == end_timestamp:
                     raise IndexError
@@ -273,9 +277,8 @@ def init_data_dataset(self):
             :U2: Wind speed (magnitude) [|m s^-1|].
             :HGT: Elevation [m].
         """
-
+        input_path = Path(Config.data_path) / "input" / Config.input_netcdf
         try:
-            input_path = os.path.join(Config.data_path, "input", Config.input_netcdf)
             self.DATA = xr.open_dataset(input_path)
         except FileNotFoundError:
             raise SystemExit(f"Input file not found at: {input_path}")

cosipy/postprocessing/field_plots/plot_cosipy_fields.py

@@ -3,6 +3,7 @@
 """
 
 import argparse
+from pathlib import Path
 import re
 
 import cartopy.crs as ccrs
@@ -409,7 +410,7 @@ def parse_arguments() -> argparse.Namespace:
         dest="file",
         required=True,
         default=None,
-        type=str,
+        type=Path,
         metavar="<path>",
         help="Path to .nc file",
     )

cosipy/postprocessing/profile_plots/plot_profiles.py

@@ -1,3 +1,7 @@
+import argparse
+import os
+from pathlib import Path
+
 import matplotlib
 # matplotlib.use("TkAgg")
 import matplotlib.pyplot as plt
@@ -242,14 +246,14 @@ def naive_fast(latvar,lonvar,lat0,lon0):
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description='Quick plot of the results file.')
-    parser.add_argument('-f', '--file', dest='file', help='Path to the result file')
+    parser.add_argument('-f', '--file', dest='file', type=Path, help='Path to the result file')
     parser.add_argument('-d', '--date', dest='pdate', help='Date of the profile plot')
     parser.add_argument('-v', '--var', dest='var', default='RHO', help='Which variable to plot (e.g. T, RHO, etc.)')
     parser.add_argument('-n', '--lat', dest='lat', default=None, help='Latitude value in case of 2D simulation', type=float)
     parser.add_argument('-m', '--lon', dest='lon', default=None, help='Longitude value in case of 2D simulation', type=float)
     parser.add_argument('-s', '--start', dest='start', default=None, help='Start date for the time plot')
     parser.add_argument('-e', '--end', dest='end', default=None, help='End date for the time plot')
-    parser.add_argument('--stake-file', dest='stake_file', default=None, help='Path to the stake data file')
+    parser.add_argument('--stake-file', dest='stake_file', default=None, type=Path, help='Path to the stake data file')
     parser.add_argument('--pit', dest='pit_name', default=None, help='Name of the pit in the stake data file')
     parser.add_argument('--depth', dest='d', nargs='+', default=None, help='An array with depth values for which the corresponding values are to be displayed', type=float)