refactor: separate printing and data handling functionality

src/pyallel/main.py

@@ -3,28 +3,81 @@
 import importlib.metadata
 import sys
 import traceback
+import time
 
 from pyallel import constants
 from pyallel.colours import Colours
 from pyallel.errors import InvalidExecutableErrors
 from pyallel.parser import Arguments, create_parser
+from pyallel.printer import Printer
 from pyallel.process_group_manager import ProcessGroupManager
 
 
-def main_loop(
-    *args: str,
-    colours: Colours,
-    interactive: bool = False,
-    timer: bool = False,
+def run_interactive(
+    process_group_manager: ProcessGroupManager, printer: Printer
 ) -> int:
-    process_group_manager = ProcessGroupManager.from_args(
-        *args,
-        colours=colours,
-        interactive=interactive,
-        timer=timer,
-    )
+    while True:
+        process_group_manager.stream()
 
-    return process_group_manager.stream()
+        printer.clear_printed_lines()
+        output = process_group_manager.get_cur_process_group_output()
+        printer.print_progress_group_output(
+            output, process_group_manager._interrupt_count
+        )
+
+        poll = process_group_manager.poll()
+        if poll is not None:
+            printer.clear_printed_lines()
+            printer.print_progress_group_output(
+                output, process_group_manager._interrupt_count, tail_output=False
+            )
+
+            if poll > 0:
+                return poll
+
+            printer.clear()
+            process_group_manager.run()
+            if not process_group_manager.next():
+                return 0
+
+        time.sleep(0.1)
+
+
+def run_non_interactive(
+    process_group_manager: ProcessGroupManager, printer: Printer
+) -> int:
+    current_process = None
+
+    while True:
+        outputs = process_group_manager.stream()
+
+        for pg in outputs.process_group_outputs.values():
+            for output in pg.processes:
+                if current_process is None:
+                    current_process = output.process
+                    output = process_group_manager.get_process(output.id)
+                    printer.print_process_output(
+                        output, include_progress=False, include_timer=False
+                    )
+                elif current_process is not output.process:
+                    continue
+                else:
+                    printer.print_process_output(output, include_cmd=False)
+
+                if output.process.poll() is not None:
+                    printer.print_process_output(output, include_output=False)
+                    current_process = None
+
+        poll = process_group_manager.poll()
+        if poll is not None:
+            if poll > 0:
+                return poll
+
+            process_group_manager.run()
+            if not process_group_manager.next():
+                return 0
+
+        time.sleep(0.1)
 
 
 def run(*args: str) -> int:
@@ -41,6 +94,7 @@ def run(*args: str) -> int:
         return 2
 
     colours = Colours.from_colour(parsed_args.colour)
+    printer = Printer(colours, timer=parsed_args.timer)
 
     interactive = True
     if not parsed_args.interactive:
@@ -50,12 +104,13 @@ def run(*args: str) -> int:
 
     message = None
     try:
-        exit_code = main_loop(
-            *parsed_args.commands,
-            colours=colours,
-            interactive=interactive,
-            timer=parsed_args.timer,
-        )
+        process_group_manager = ProcessGroupManager.from_args(*parsed_args.commands)
+        process_group_manager.run()
+
+        if interactive:
+            exit_code = run_interactive(process_group_manager, printer)
+        else:
+            exit_code = run_non_interactive(process_group_manager, printer)
     except InvalidExecutableErrors as e:
         exit_code = 1
         message = str(e)
@@ -65,17 +120,11 @@ def run(*args: str) -> int:
 
     if exit_code == 1:
         if not message:
-            print(
-                f"{colours.dim_on}=>{colours.dim_off} {colours.red_bold}Failed!{colours.reset_colour}"
-            )
+            printer.error("\nFailed!")
         else:
-            print(
-                f"{colours.dim_on}=>{colours.dim_off} {colours.red_bold}Error: {message}{colours.reset_colour}"
-            )
+            printer.error(f"Error: {message}")
     elif exit_code == 0:
-        print(
-            f"{colours.dim_on}=>{colours.dim_off} {colours.green_bold}Done!{colours.reset_colour}"
-        )
+        printer.ok("\nDone!")
 
     return exit_code