update run-fit so that you can ctrl-c it

1 files changed, 16 insertions, 1 deletions

diff --git a/src/run-fit b/src/run-fit
index 6c8bb58..203aa59 100755
--- a/src/run-fit
+++ b/src/run-fit
@@ -13,6 +13,8 @@ def run_fit(filename):
 if __name__ == '__main__':
     import argparse
     from multiprocessing import Pool, cpu_count
+    import signal
+    import os
 
     parser = argparse.ArgumentParser("fit multiple zdab files")
     parser.add_argument("-j", "--jobs", type=int, default=None,
@@ -26,5 +28,18 @@ if __name__ == '__main__':
     if jobs is None:
         jobs = cpu_count()
 
+    # see https://stackoverflow.com/questions/11312525/catch-ctrlc-sigint-and-exit-multiprocesses-gracefully-in-python
+    handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
     p = Pool(jobs)
-    p.map(run_fit,args.filenames)
+    signal.signal(signal.SIGINT, handler)
+    try:
+        result = p.map_async(run_fit,args.filenames)
+        result.get()
+    except KeyboardInterrupt:
+        print("ctrl-c caught. quitting...")
+        p.terminate()
+        os.killpg(os.getpgid(0),9)
+    else:
+        p.close()
+
+    p.join()