Adds progress bar which uses the progress event (#62)

Author: StephenNneji

RATapi/run.py

@@ -1,11 +1,53 @@
 import time
 
+from tqdm import tqdm
+
 import RATapi.rat_core
 from RATapi.inputs import make_input
 from RATapi.outputs import make_results
 from RATapi.utils.enums import Display
 
 
+class ProgressBar:
+    """Creates a progress bar that gets updates from the progress event during a
+    calculation
+
+    Parameters
+    ----------
+    display : bool, default: True
+            Indicates if displaying is allowed
+
+    """
+
+    def __init__(self, display=True):
+        self.display = display
+
+    def __enter__(self):
+        if self.display:
+            RATapi.events.register(RATapi.events.EventTypes.Progress, self.updateProgress)
+        self.pbar = tqdm(total=100, desc="", delay=1, bar_format="{l_bar}{bar}", ncols=90, disable=not self.display)
+        self.pbar.delay = 0
+        return self.pbar
+
+    def updateProgress(self, event):
+        """Callback for the progress event.
+
+        Parameters
+        ----------
+        event: ProgressEventData
+            The progress event data.
+        """
+
+        value = event.percent * 100
+        self.pbar.desc = event.message
+        self.pbar.update(value - self.pbar.n)
+
+    def __exit__(self, _exc_type, _exc_val, _traceback):
+        self.pbar.leave = False
+        if self.display:
+            RATapi.events.clear(RATapi.events.EventTypes.Progress, self.updateProgress)
+
+
 def run(project, controls):
     """Run RAT for the given project and controls inputs."""
     parameter_field = {
@@ -19,23 +61,24 @@ def run(project, controls):
     }
 
     horizontal_line = "\u2500" * 107 + "\n"
-
+    display_on = controls.display != Display.Off
     problem_definition, cells, limits, priors, cpp_controls = make_input(project, controls)
 
-    if controls.display != Display.Off:
+    if display_on:
         print("Starting RAT " + horizontal_line)
 
     start = time.time()
-    problem_definition, output_results, bayes_results = RATapi.rat_core.RATMain(
-        problem_definition,
-        cells,
-        limits,
-        cpp_controls,
-        priors,
-    )
+    with ProgressBar(display=display_on):
+        problem_definition, output_results, bayes_results = RATapi.rat_core.RATMain(
+            problem_definition,
+            cells,
+            limits,
+            cpp_controls,
+            priors,
+        )
     end = time.time()
 
-    if controls.display != Display.Off:
+    if display_on:
         print(f"Elapsed time is {end-start:.3f} seconds\n")
 
     results = make_results(controls.procedure, output_results, bayes_results)
@@ -45,7 +88,7 @@ def run(project, controls):
         for index, value in enumerate(getattr(problem_definition, parameter_field[class_list])):
             getattr(project, class_list)[index].value = value
 
-    if controls.display != Display.Off:
+    if display_on:
         print("Finished RAT " + horizontal_line)
 
     return project, results

requirements.txt

@@ -7,4 +7,5 @@ pytest-cov >= 4.1.0
 matplotlib >= 3.8.3
 StrEnum >= 0.4.15; python_version < '3.11'
 ruff >= 0.4.10
-scipy >= 1.13.1
+scipy >= 1.13.1
+tqdm >= 4.66.5

setup.py

@@ -171,6 +171,7 @@ def build_libraries(self, libraries):
         "pydantic >= 2.7.2",
         "matplotlib >= 3.8.3",
         "scipy >= 1.13.1",
+        "tqdm>=4.66.5",
     ],
     extras_require={
         ':python_version < "3.11"': ["StrEnum >= 0.4.15"],

tests/test_run.py

@@ -11,6 +11,8 @@
 import RATapi
 import RATapi.outputs
 import RATapi.rat_core
+from RATapi.events import EventTypes, ProgressEventData, notify
+from RATapi.run import ProgressBar
 from RATapi.utils.enums import Calculations, Geometries, LayerModels, Procedures
 from tests.utils import check_results_equal
 
@@ -320,3 +322,20 @@ def test_run(test_procedure, test_output_problem, test_output_results, test_baye
         project, results = RATapi.run(input_project, RATapi.Controls(procedure=test_procedure))
 
     check_results_equal(test_results, results)
+
+
+def test_progress_bar() -> None:
+    event = ProgressEventData()
+    event.message = "TESTING"
+    event.percent = 0.2
+    with ProgressBar() as bar:
+        assert bar.n == 0
+        notify(EventTypes.Progress, event)
+        assert bar.desc == "TESTING"
+        assert bar.n == 20
+
+        event.message = "AGAIN"
+        event.percent = 0.6
+        notify(EventTypes.Progress, event)
+        assert bar.desc == "AGAIN"
+        assert bar.n == 60