Merge pull request #98 from seibert-media/little-refactorings

parallel_runner: Use defaultdict
Fix: Crashes when switching screens in parallel mode

Author: jpaul-seibert

automatix/parallel_runner.py

@@ -2,6 +2,7 @@
 import pickle
 import subprocess
 from argparse import Namespace
+from collections import defaultdict
 from tempfile import TemporaryDirectory
 from time import time
 
@@ -12,12 +13,10 @@
 
 
 def get_batch_groups(batch_items: list) -> dict:
-    batch_groups = {'_default_': []}
+    batch_groups = defaultdict(list)
     for batch_item in batch_items:
         if group := batch_item.get('group'):
             assert group != '_default_', 'The group name "_default_" is reserved. Please use something different.'
-            if group not in batch_groups.keys():
-                batch_groups[group] = []
             batch_groups[group].append(batch_item)
         else:
             batch_groups['_default_'].append(batch_item)
@@ -38,7 +37,7 @@ def create_auto_files(script: dict, batch_items: list, args: Namespace, tempdir:
     LOG.info(f'Using temporary directory to save object files: {tempdir}')
 
     batch_groups = get_batch_groups(batch_items=batch_items)
-    default_group = batch_groups.pop('_default_')
+    default_group = batch_groups.pop('_default_', [])
 
     digits = len(str(len(batch_groups) + len(default_group)))
 

automatix/parallel_ui.py

@@ -168,7 +168,7 @@ def process_user_input(cw: CursesWriter, autos: Autos) -> str | None:
         cw.input_buffer += char
 
 
-def parallel_ui(stdscr: curses.window, tempdir: str):
+def parallel_ui(stdscr: curses.window, tempdir: str) -> tuple[str, str | None]:
     cw = CursesWriter(stdscr=stdscr)
 
     # Wait until the status file exists
@@ -194,24 +194,23 @@ def parallel_ui(stdscr: curses.window, tempdir: str):
             cw.stdscr.nodelay(False)
             while cw.stdscr.getch() not in [ord('q'), ord('Q')]:
                 sleep(0.1)
-            return 'quit'
+            return 'quit', None
 
         # 4. Process user input
         screen_to_switch = process_user_input(cw=cw, autos=autos)
         if screen_to_switch:
-            # Important: End curses before an external program takes control of the terminal
-            curses.endwin()
-            print(f"Switching to screen '{screen_to_switch}'... (Return with <ctrl>+a d)")
-            subprocess.run(['screen', '-r', screen_to_switch])
-            return 'restart'
+            return 'restart', screen_to_switch
 
         sleep(0.2)  # Short pause to reduce CPU load
 
 
 def screen_switch_loop(tempdir: str):
     while True:
         try:
-            exit_reason = curses.wrapper(parallel_ui, tempdir)
+            exit_reason, screen_to_switch = curses.wrapper(parallel_ui, tempdir)
+            if screen_to_switch:
+                print(f"Switching to screen '{screen_to_switch}'... (Return with <ctrl>+a d)")
+                subprocess.run(['screen', '-r', screen_to_switch])
         except curses.error as exc:
             LOG.error(f"Curses error: {exc}")
             LOG.error("Could not start the curses interface. Is the terminal compatible?")