From c83d2d9bdbd00e4a7594682a65d4275cecbcdd70 Mon Sep 17 00:00:00 2001
From: Sven Augustin <sven.augustin@psi.ch>
Date: Mon, 14 Mar 2022 18:14:36 +0100
Subject: [PATCH] made condition stoppable; added some colored printing

---
 slic/core/condition/condition.py | 17 +++++++++++++++--
 slic/core/scanner/scanbackend.py | 18 ++++++++++++++++--
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/slic/core/condition/condition.py b/slic/core/condition/condition.py
index 373e7cb8..d3e9c572 100644
--- a/slic/core/condition/condition.py
+++ b/slic/core/condition/condition.py
@@ -20,6 +20,7 @@ class Condition(BaseCondition):
         self._repeater_gen = self._repeater()
 
         self.data = []
+        self.running = False
 
 
     def check(self):
@@ -31,7 +32,11 @@ class Condition(BaseCondition):
         raise NotImplementedError
 
     def sleep(self):
-        sleep(self.wait_time)
+        time_start = time()
+        while time() - time_start < self.wait_time:
+            if not self.running:
+                return
+            sleep(0.01)
 
 
     def clear_and_start_counting(self):
@@ -47,6 +52,9 @@ class Condition(BaseCondition):
     def start_counting(self):
         raise NotImplementedError
 
+    def stop(self):
+        self.running = False
+        self.stop_counting()
 
     def stop_counting_and_analyze(self):
         self.stop_counting()
@@ -91,8 +99,11 @@ class Condition(BaseCondition):
     def get_ready(self):
         time_start = time()
         was_ever_unhappy = False
+        self.running = True
 
         while not self.long_check():
+            if not self.running:
+                break
             was_ever_unhappy = True
             delta_t = time() - time_start
             print(f"Condition is unhappy, waiting for OK conditions since {delta_t:.1f} seconds.")
@@ -131,9 +142,11 @@ class Condition(BaseCondition):
             while True:
                 self.get_ready()
                 yield True # signal condition wants a repeat (incl. the first measurement)
+                if not self.running:
+                    break
                 if self.is_happy():
                     break
-            yield False # signal condition was happy with last repeat
+            yield False # signal condition was happy with last repeat or stopped
 
 
 
diff --git a/slic/core/scanner/scanbackend.py b/slic/core/scanner/scanbackend.py
index 5e0f3fb2..b4cf0435 100644
--- a/slic/core/scanner/scanbackend.py
+++ b/slic/core/scanner/scanbackend.py
@@ -1,4 +1,5 @@
 import os
+import colorama
 
 from slic.utils import make_missing_dir, printable_exception
 from slic.utils.printing import printable_dict, itemize, format_header, printable_table
@@ -132,16 +133,27 @@ class ScanBackend:
         ntotal = len(values)
         for n, val in enumerate(values):
             if not self.running:
+                n -= 1 # stopped before this iteration
                 break
+            print(colorama.Fore.GREEN, end="")
             print("Scan step {} of {}".format(n+1, ntotal))
+            print(colorama.Fore.RESET, end="")
             do_step(n, val, step_info=step_info)
 
-        print("All scan steps done")
+        if self.running and n+1 == ntotal:
+            print(colorama.Fore.GREEN, end="")
+            print("All scan steps done")
+            print(colorama.Fore.RESET, end="")
+        else:
+            print(colorama.Fore.RED, end="")
+            print("Stopped during scan step {} of {}".format(n+1, ntotal))
+            print(colorama.Fore.RESET, end="")
 
 
     def do_checked_step(self, *args, **kwargs):
         while self.condition.wants_repeat():
-            self.do_step(*args, **kwargs)
+            if self.running:
+                self.do_step(*args, **kwargs)
 
 
     def do_step(self, n_step, step_values, step_info=None):
@@ -234,6 +246,8 @@ class ScanBackend:
     def stop(self):
         self.running = False
         stop_all(self.current_tasks)
+        if self.condition:
+            self.condition.stop()
 
 
     def __repr__(self):