Refactor clean_rundir and related

Author: bernstei

tests/local_scripts/complete_pytest.tin

@@ -29,7 +29,7 @@ echo "" >> complete_pytest.tin.out
 # buildcell
 export WFL_PYTEST_BUILDCELL=$HOME/src/work/AIRSS/airss-0.9.1/src/buildcell/src/buildcell
 # VASP
-module load dft/vasp
+module load dft/vasp/serial
 export ASE_VASP_COMMAND=vasp.serial
 export ASE_VASP_COMMAND_GAMMA=vasp.gamma.serial
 export PYTEST_VASP_POTCAR_DIR=$VASP_PATH/pot/rev_54/PBE

wfl/calculators/castep.py

@@ -25,7 +25,7 @@ class Castep(WFLFileIOCalculator, ASE_Castep):
         what kind of files to keep from the run
             - True : everything kept
             - None, False : nothing kept, unless calculation fails
-            - "default"   : only ones needed for NOMAD uploads ('\*.pwo')
+            - "default"   : only ones needed for NOMAD uploads ('\*.{castep,param,cell}')
             - list(str)   : list of file globs to save
     rundir_prefix: str / Path, default 'run\_CASTEP\_'
         Run directory name prefix

wfl/calculators/orca/__init__.py

@@ -34,7 +34,7 @@ class ORCA(WFLFileIOCalculator, ASE_ORCA):
         what kind of files to keep from the run
             - True : everything kept
             - None, False : nothing kept, unless calculation fails
-            - "default"   : only ones needed for NOMAD uploads ('\*.pwo')
+            - "default"   : only ones needed for NOMAD uploads ('\*.{inp,out,engrad,xyz,_traj.xyz}')
             - list(str)   : list of file globs to save
     rundir_prefix: str / Path, default 'ORCA\_'
         Run directory name prefix

wfl/calculators/utils.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from wfl.utils.file_utils import clean_dir
 
 def clean_rundir(rundir, keep_files, default_keep_files, calculation_succeeded):
@@ -7,19 +9,30 @@ def clean_rundir(rundir, keep_files, default_keep_files, calculation_succeeded):
     ----------
     rundir: str
         path to run dir
-    keep_files: 'default' / list(str) / '*' / bool / None
+    keep_files: 'default' / iterable(file_glob) / '*' / bool / None
         files to keep, None or False for nothing, '*' or True for all
     default_keep_files: list(str)
         files to keep if keep_files == 'default' or calculation_succeeded is False
+        and keep_files is not True
     calculation_succeeded: bool
     """
-    if not calculation_succeeded:
-        clean_dir(rundir, set(default_keep_files) | set(keep_files if keep_files else []), force=False)
-    elif keep_files == 'default':
-        clean_dir(rundir, default_keep_files, force=False)
-    elif not keep_files:
-        clean_dir(rundir, False, force=False)
-    else:
-        clean_dir(rundir, keep_files, force=False)
+    if isinstance(keep_files, str):
+        if keep_files == 'default':
+            keep_files = default_keep_files
+        elif keep_files == '*'
+            keep_files = [keep_files]
+        else:
+            raise ValueError(f"str keep_files can only be 'default' or '*', not '{keep_files}'")
 
+    # now keep_files should either be True, evaluate to False, or an iterable
+
+    # calculation failed, default optionally union with keep_files
+    if not calculation_succeeded:
+        if not keep_files:
+            keep_files = default_keep_files
+        elif keep_files is not True:
+            keep_files = set(default_keep_files) | set(keep_files)
+        else:
+            raise ValueError(f"Unknown type or value for keep_files '{keep_files}'")
 
+    clean_dir(rundir, keep_files, force=False)

wfl/utils/file_utils.py

@@ -9,11 +9,11 @@ def clean_dir(directory, keep_files, force=False):
     Parameters
     ----------
     directory : directory to be cleaned
-    keep_files: bool or list(filenames) or str
+    keep_files: bool or iterable(file_globs)
         What to keep in rundir when done:
-            - list(filenames) : ONLY these filenames if they exist
-            - True or "*" : everything - does nothing
-            - False or None : remove everything, or anything evaluating to False in if
+            - True : keep everything
+            - iterable : ONLY files matching these globs if they exist
+            - evaluates to False : nothing
     force : bool, default = False
         fail if directory does not exist
 
@@ -29,21 +29,19 @@ def clean_dir(directory, keep_files, force=False):
         else:
             return
 
-    # defaults
-    if keep_files is None:
-        keep_files = False
-    elif keep_files == "*":
-        keep_files = True
-    elif isinstance(keep_files, str):
-        keep_files = [keep_files]
+    if isinstance(keep_files, str):
+        # common error that will not be caught below
+        raise ValueError(f"keep_files must not be str '{keep_files}'")
 
     # operations
-    if isinstance(keep_files, bool) and keep_files:
+    if keep_files is True:
+        # bool True, keep all
         return
     elif not keep_files:
-        # lets None and anything else evaluating to False
+        # evaluates to False, keep none
         shutil.rmtree(directory)
-    elif isinstance(keep_files, (list, tuple, set)):
+    else:
+        # keep_files should be an iterable
         extended_keep_files = []
         for glob_index in keep_files:
             extended_keep_files.extend(glob(os.path.join(directory, glob_index)))
@@ -56,5 +54,3 @@ def clean_dir(directory, keep_files, force=False):
                     os.remove(abs_fn)
                 else:
                     shutil.rmtree(abs_fn)
-    else:
-        raise RuntimeError('Got unknown type or value for keep_rundir \'{}\''.format(keep_files))