Merge pull request #30 from molssi-seamm/dev

Enhancement: loop control

HISTORY.rst

@@ -1,6 +1,11 @@
 =======
 History
 =======
+2024.11.14: Enhancement: loop control
+   * Added methods in the loop_step module to break from a loop, break_loop(), continue
+     a loop, continue_loop(), and totallly skip an iteration, removing any files
+     assocatied with it, skip_iteration.
+
 2024.10.15: Bugfix: incorrectly wrapping text description of substeps
    * When printing the description of the substeps in the loop, the code incorrectly
      wrapped the text, which caused issues with e.g. tables.

loop_step/__init__.py

@@ -8,7 +8,15 @@
 # Bring up the classes so that they appear to be directly in
 # the loop_step package.
 
-from loop_step.loop import Loop  # noqa: F401
+from loop_step.loop import (  # noqa: F401
+    Loop,
+    BreakLoop,
+    break_loop,
+    ContinueLoop,
+    continue_loop,
+    SkipIteration,
+    skip_iteration,
+)
 from loop_step.loop_parameters import LoopParameters  # noqa: F401
 from loop_step.loop_step import LoopStep  # noqa: F401
 from loop_step.tk_loop import TkLoop  # noqa: F401

loop_step/loop.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 import re
 import shlex
+import shutil
 import sys
 import traceback
 
@@ -24,6 +25,43 @@
 printer = printing.getPrinter("loop")
 
 
+class BreakLoop(Exception):
+    """Indicates that SEAMM should break from the loop"""
+
+    def __init__(self, message="break from the loop"):
+        super().__init__(message)
+
+
+def break_loop():
+    """Break from the loop and continue on."""
+    raise BreakLoop()
+
+
+class ContinueLoop(Exception):
+    """Indicates that SEAMM should continue from the loop"""
+
+    def __init__(self, message="continue with next iteration of loop"):
+        super().__init__(message)
+
+
+def continue_loop():
+    """Continue to the next iteration of the loop"""
+    raise ContinueLoop()
+
+
+class SkipIteration(Exception):
+    """Indicates that SEAMM should skip this iteration of the loop,
+    removing any directories, etc."""
+
+    def __init__(self, message="skip iteration of loop"):
+        super().__init__(message)
+
+
+def skip_iteration():
+    """Entirely skip this iteration, removing any files, etc."""
+    raise SkipIteration()
+
+
 class Loop(seamm.Node):
     def __init__(self, flowchart=None, extension=None):
         """Setup the non-graphical part of the Loop step in a
@@ -688,6 +726,13 @@ def run(self):
                 printer.normal("\nDeprecation warning: " + str(e))
                 traceback.print_exc(file=sys.stderr)
                 traceback.print_exc(file=sys.stdout)
+            except BreakLoop:
+                break
+            except ContinueLoop:
+                next_node = self
+            except SkipIteration:
+                next_node = self
+                shutil.rmtree(iter_dir)
             except Exception as e:
                 tmp = self.working_path.name
                 printer.job(f"Caught exception in loop iteration {tmp}: {str(e)}")