Add option to specify output name

* Reorganise to have common stuff in a single place
* add iterator to further remove code duplication
* Add support for reading "."

xemc3/cli/_common.py

@@ -0,0 +1,43 @@
+import os
+from argparse import ArgumentParser
+
+
+def parser(desc, ext="nc") -> ArgumentParser:
+    parser = ArgumentParser(description=desc)
+    parser.add_argument(
+        "path",
+        nargs="+",
+        help=f"""Path of the directory to load. The netcdf file will be
+        called dir.{ext} if the folder was called dir.""",
+    )
+    parser.add_argument("-q", "--quiet", action="store_true", help="Be less verbose")
+    parser.add_argument(
+        "-o",
+        "--name",
+        help=f"Specify the name for the output file. Defaults to `dir.{ext}` "
+        f"when not given. Otherwise `name.{ext}` is used.",
+    )
+    return parser
+
+
+def iter_dir(args):
+    if len(args.path) > 1 and args.name:
+        raise RuntimeError(
+            """Providing an explicit output name and more than one folder to be read is
+currently not supported. Either use the default output name or read one file
+after another. Further you can use the python interface that gives more
+control."""
+        )
+    for d in args.path:
+        while d[-1] == "/":
+            d = d[:-1]
+        if args.name:
+            yield d, args.name
+        else:
+            if os.path.isdir(d):
+                if d == ".":
+                    yield d, os.getcwd()
+                else:
+                    yield d, d
+            else:
+                yield d, d.rsplit(".", 1)[0]

xemc3/cli/append_time.py

@@ -1,67 +1,52 @@
 #!/usr/bin/env python3
-
 import os
 import sys
-from argparse import ArgumentParser
 
 import xarray as xr
 
 from ..core import dataset
 from ..core.load import load_all
+from ._common import iter_dir, parser
 
 
-def append_time(d: str, verbose: bool = False) -> None:
-    while d[-1] == "/":
-        d = d[:-1]
+def append_time(read: str, name: str, verbose: bool = False) -> None:
     try:
-        os.replace(d + ".nc", d + ".old.nc")
-        old = xr.open_dataset(d + ".old.nc")
+        os.replace(name + ".nc", name + ".old.nc")
+        old = xr.open_dataset(name + ".old.nc")
     except FileNotFoundError:
         old = None
     try:
         if verbose:
-            print(f"Loading {d} ...", end="")
+            print(f"Loading {read} ...", end="")
             sys.stdout.flush()
-        ds = load_all(d)
+        ds = load_all(read)
         if old is not None:
             ds = xr.concat([old, ds], "time", "different")
         if verbose:
             print(" writing ...", end="")
             sys.stdout.flush()
-        ds.emc3.to_netcdf(d + ".nc")
+        ds.emc3.to_netcdf(name + ".nc")
     except Exception:
         if old is not None:
-            os.replace(d + ".old.nc", d + ".nc")
+            os.replace(name + ".old.nc", name + ".nc")
         raise
     if old is not None:
-        os.remove(d + ".old.nc")
+        os.remove(name + ".old.nc")
     if verbose:
         print(" done")
 
 
-def parser() -> ArgumentParser:
-    parser = ArgumentParser(
-        description="""
+def main() -> None:
+    args = parser(
+        """
         Load the data from EMC3 simulations and store as netcdf file. The
         data is appended for each simulation to the netcdf file, which
         will be created if it does not yet exists.
         """
-    )
-    parser.add_argument(
-        "path",
-        nargs="+",
-        help="""Path of the directory to load. The netcdf file will be
-        called dir.nc if the folder was called dir.""",
-    )
-    parser.add_argument("-q", "--quiet", action="store_true", help="Be less verbose")
-    return parser
-
-
-def main() -> None:
-    args = parser().parse_args()
+    ).parse_args()
 
-    for d in args.path:
-        append_time(d, not args.quiet)
+    for d, n in iter_dir(args):
+        append_time(d, n, not args.quiet)
 
 
 if __name__ == "__main__":

xemc3/cli/to_archive.py

@@ -7,45 +7,41 @@
 import xarray as xr
 
 from ..core.load import archive, load_all
+from ._common import iter_dir
+from ._common import parser as commonparser
 
+ext = "arch.nc"
 
-def to_archive(d: str, quiet: bool, geom: bool, mapping: bool, delete: bool) -> None:
-    while d[-1] == "/":
-        d = d[:-1]
+
+def to_archive(
+    read: str, name: str, quiet: bool, geom: bool, mapping: bool, delete: bool
+) -> None:
     fromds = False
     if not quiet:
-        print(f"Loading {d} ...", end="")
+        print(f"Loading {read} ...", end="")
         sys.stdout.flush()
-    if os.path.isdir(d):
-        ds = load_all(d)
+    if os.path.isdir(read):
+        ds = load_all(read)
     else:
-        ds = xr.open_dataset(d)
-        dorg = d
-        d = ".".join(d.split(".")[:-1])
+        ds = xr.open_dataset(read)
         fromds = True
     if not quiet:
         print(" writing ...", end="")
         sys.stdout.flush()
-    archive(ds, d + ".arch.nc", geom, mapping)
+    archive(ds, f"{name}.{ext}", geom, mapping)
     if delete and fromds:
         if not quiet:
             print(" deleting ...", end="")
             sys.stdout.flush()
-        os.unlink(dorg)
+        os.unlink(read)
     if not quiet:
         print(" done")
 
 
 def parser() -> ArgumentParser:
-    parser = ArgumentParser(
-        description="Archive the data from EMC3 simulations as a netcdf file."
-    )
-    parser.add_argument(
-        "path",
-        nargs="+",
-        help="Path of the directory to load or a netcdf file. The netcdf file will be called dir.arc.nc if the folder was called dir.",
+    parser = commonparser(
+        "Archive the data from EMC3 simulations as a netcdf file.", ext=ext
     )
-    parser.add_argument("-q", "--quiet", action="store_true", help="Be less verbose")
     parser.add_argument(
         "-g", "--geometry", action="store_true", help="Also store geometry"
     )
@@ -59,17 +55,19 @@ def parser() -> ArgumentParser:
         "-d",
         "--delete",
         action="store_true",
-        help="Delete the uncompressed input file. Only used if the input was a netcdf file",
+        help="Delete the uncompressed input file. "
+        "Only used if the input was a netcdf file",
     )
     return parser
 
 
 def main() -> None:
     args = parser().parse_args()
 
-    for d in args.path:
+    for d, n in iter_dir(args):
         to_archive(
             d,
+            n,
             args.quiet,
             geom=args.geometry,
             mapping=not args.no_mapping,

xemc3/cli/to_netcdf.py

@@ -1,46 +1,33 @@
 #!/usr/bin/env python3
 
 import sys
-from argparse import ArgumentParser
 
 from ..core import dataset
 from ..core.load import load_all
+from ._common import iter_dir, parser
 
 
-def to_netcdf(d: str, quiet: bool = True) -> None:
-    while d[-1] == "/":
-        d = d[:-1]
+def to_netcdf(read: str, name: str, quiet: bool = True) -> None:
     if not quiet:
-        print(f"Loading {d} ...", end="")
+        print(f"Loading {read} ...", end="")
         sys.stdout.flush()
-    ds = load_all(d)
+    ds = load_all(read)
     if not quiet:
         print(" writing ...", end="")
         sys.stdout.flush()
-    ds.emc3.to_netcdf(d + ".nc")
+    ds.emc3.to_netcdf(name + ".nc")
     if not quiet:
         print(" done")
 
 
-def parser() -> ArgumentParser:
-    parser = ArgumentParser(
-        description="Load the data from EMC3 simulations and store as netcdf file. "
-        "The data is written for each simulation to a netcdf file."
-    )
-    parser.add_argument(
-        "path",
-        nargs="+",
-        help="Path of the directory to load. The netcdf file will be called dir.nc if the folder was called dir.",
-    )
-    parser.add_argument("-q", "--quiet", action="store_true", help="Be less verbose")
-    return parser
-
-
 def main() -> None:
-    args = parser().parse_args()
+    args = parser(
+        "Load the data from EMC3 simulations and store as netcdf file. "
+        "The data is written for each simulation to a netcdf file."
+    ).parse_args()
 
-    for d in args.path:
-        to_netcdf(d, args.quiet)
+    for d, n in iter_dir(args):
+        to_netcdf(d, n, args.quiet)
 
 
 if __name__ == "__main__":