diff --git a/xemc3/cli/_common.py b/xemc3/cli/_common.py new file mode 100644 index 0000000..7a8f1c2 --- /dev/null +++ b/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] diff --git a/xemc3/cli/append_time.py b/xemc3/cli/append_time.py index 945af2a..2eba7ff 100755 --- a/xemc3/cli/append_time.py +++ b/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__": diff --git a/xemc3/cli/to_archive.py b/xemc3/cli/to_archive.py index 45c67de..37476d5 100644 --- a/xemc3/cli/to_archive.py +++ b/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,7 +55,8 @@ 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 @@ -67,9 +64,10 @@ def parser() -> ArgumentParser: 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, diff --git a/xemc3/cli/to_netcdf.py b/xemc3/cli/to_netcdf.py index 7e6e075..11f02ea 100755 --- a/xemc3/cli/to_netcdf.py +++ b/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__":