fix: allocate more lines to processes with more output

src/pyallel/printer.py

@@ -259,46 +259,52 @@ def set_process_lines(
         lines -= 2
 
     # Allocate lines to processes that have a fixed percentage of lines set
+    allocated_process_lines = lines // len(output.processes)
+    processes_with_dynamic_lines: list[ProcessOutput] = []
     used_lines = 0
-    process_with_most_lines: ProcessOutput | None = None
-    other_processes: list[ProcessOutput] = []
-    for out in output.processes:
-        if not out.process.percentage_lines:
-            other_processes.append(out)
+    for process_output in output.processes:
+        # This process output doesn't have percentage_lines set, so skip it
+        if not process_output.process.percentage_lines:
+            processes_with_dynamic_lines.append(process_output)
             continue
 
-        out.process.lines = int(lines * out.process.percentage_lines)
-        used_lines += out.process.lines
-
-        if (
-            process_with_most_lines is None
-            or process_with_most_lines.process.lines < out.process.lines
-        ):
-            process_with_most_lines = out
+        process_output.process.lines = int(
+            lines * process_output.process.percentage_lines
+        )
+        used_lines += process_output.process.lines
 
+    # Remove the used lines from the total available lines
     lines -= used_lines
 
-    # Allocate the rest of the available lines to the other processes that don't have fixed lines set
-    num_dynamic_processes = len(other_processes)
-    if num_dynamic_processes:
-        remainder = lines % num_dynamic_processes
-        tail = lines // num_dynamic_processes
-
-        for out in other_processes:
-            out.process.lines = tail
-
-        # If we have lines left to allocate, we give all of them to the process with the highest 
-        # allocated lines or to the last process
-        if remainder:
-            if process_with_most_lines:
-                process_with_most_lines.process.lines += remainder
-            else:
-                other_processes[-1].process.lines += remainder
-    else:
-        # Otherwise allocate remaining lines to the process with the highest allocated lines
-        remainder = lines
-        if remainder and process_with_most_lines:
-            process_with_most_lines.process.lines += remainder
+    while lines:
+        # Calculate how many lines each process should have based on how many processes and lines are left
+        num_processes = len(processes_with_dynamic_lines) or 1
+        allocated_process_lines = lines // num_processes
+        processes_with_excess_output: list[ProcessOutput] = []
+        recalculate_lines = False
+        for process_output in processes_with_dynamic_lines:
+            # If the number of lines in this process output is less than how many terminal lines we would allocate it,
+            # Set it's allocated terminal lines to the exact number of lines in its output and remove this number from
+            # the total available terminal lines
+            if process_output.lines < allocated_process_lines:
+                process_output.process.lines = process_output.lines
+                lines -= process_output.process.lines
+                recalculate_lines = True
+                continue
+
+            processes_with_excess_output.append(process_output)
+
+        # We need to re-calcuate how many terminal lines we can allocate to each process if the output of at least one process
+        # contains less lines than what we would normally allocate it. This is done so we can allocate these extra lines to the
+        # other processes that contain more lines of output.
+        if recalculate_lines:
+            processes_with_dynamic_lines = processes_with_excess_output
+        else:
+            # All remaining processes exceed the number of terminal lines we will allocate them, so allocate them
+            # their terminal lines as normal and break out of the while loop
+            for process_output in processes_with_excess_output:
+                process_output.process.lines = allocated_process_lines
+            break
 
 
 def get_num_lines(line: str, columns: int | None = None) -> int:

src/pyallel/process.py

@@ -13,10 +13,12 @@ class ProcessOutput:
     def __init__(self, id: int, process: Process, data: str = "") -> None:
         self.id = id
         self.data = data
+        self.lines = len(data.splitlines()) + 1
         self.process = process
 
     def merge(self, other: ProcessOutput) -> None:
         self.data += other.data
+        self.lines += len(other.data.splitlines())
 
 
 class Process:

tests/test_printer.py

@@ -53,46 +53,58 @@ def test_get_num_lines_ignores_ansi_chars(chars: str) -> None:
 def test_set_process_lines() -> None:
     output = ProcessGroupOutput(
         id=1,
-        processes=[ProcessOutput(id=1, process=Process(1, "echo first; echo second"))],
+        processes=[
+            ProcessOutput(
+                id=1,
+                process=Process(1, "echo first; echo second"),
+                data="first\nsecond\n",
+            )
+        ],
     )
 
     set_process_lines(output, lines=58)
 
-    assert output.processes[0].process.lines == 58
+    assert output.processes[0].process.lines == 3
 
 
-@pytest.mark.parametrize(
-    "lines,expected_lines1,expected_lines2,expected_lines3",
-    (
-        pytest.param(59, 19, 19, 21, id="59 lines shared between 3 processes"),
-        pytest.param(50, 16, 16, 18, id="50 lines shared between 3 processes"),
-        pytest.param(31, 10, 10, 11, id="31 lines shared between 3 processes"),
-    ),
-)
-def test_set_process_lines_shares_lines_across_processes(
-    lines: int, expected_lines1: int, expected_lines2: int, expected_lines3: int
-) -> None:
+def test_set_process_lines_shares_lines_across_processes() -> None:
     output = ProcessGroupOutput(
         id=1,
         processes=[
-            ProcessOutput(id=1, process=Process(1, "echo first; echo second")),
-            ProcessOutput(id=2, process=Process(2, "echo first; echo second")),
-            ProcessOutput(id=3, process=Process(3, "echo first; echo second")),
+            ProcessOutput(
+                id=1,
+                process=Process(1, "echo first; echo second"),
+                data="first\nsecond\n",
+            ),
+            ProcessOutput(
+                id=2,
+                process=Process(2, "echo first; echo second"),
+                data="first\nsecond\n",
+            ),
+            ProcessOutput(
+                id=3,
+                process=Process(3, "echo first; echo second"),
+                data="first\nsecond\n",
+            ),
         ],
     )
 
-    set_process_lines(output, lines=lines)
+    set_process_lines(output, lines=59)
 
-    assert output.processes[0].process.lines == expected_lines1
-    assert output.processes[1].process.lines == expected_lines2
-    assert output.processes[2].process.lines == expected_lines3
+    assert output.processes[0].process.lines == 3
+    assert output.processes[1].process.lines == 3
+    assert output.processes[2].process.lines == 3
 
 
 def test_set_process_lines_shares_lines_across_many_more_processes() -> None:
     output = ProcessGroupOutput(
         id=1,
         processes=[
-            ProcessOutput(id=i, process=Process(i, "echo first; echo second"))
+            ProcessOutput(
+                id=i,
+                process=Process(i, "echo first; echo second"),
+                data="first\nsecond\n",
+            )
             for i in range(1, 60)
         ],
     )
@@ -107,29 +119,12 @@ def test_set_process_lines_shares_lines_across_many_more_processes() -> None:
     "lines,lines1,lines2,lines3,expected_lines1,expected_lines2,expected_lines3",
     (
         pytest.param(
-            59,
-            0.4,
-            0.2,
-            0.2,
-            37,
-            11,
-            11,
-            id="59 lines shared between 3 processes with remainder given to process with most lines",
-        ),
-        pytest.param(
-            59,
-            0.4,
-            0.2,
-            0.0,
-            23,
-            11,
-            25,
-            id="59 lines shared between 3 processes with remainder given to last process",
+            59, 0.4, 0.2, 0.2, 23, 11, 11, id="59 lines shared between 3 processes"
         ),
-        pytest.param(59, 1.0, 0.0, 0.0, 59, 0, 0, id="59 lines given to first process"),
         pytest.param(
-            59, 0.5, 0.0, 0.0, 29, 15, 15, id="59 lines given to first process"
+            59, 1.0, 0.0, 0.0, 59, 0, 0, id="All lines given to first process"
         ),
+        pytest.param(59, 0.5, 0.0, 0.0, 29, 3, 3, id="29 lines given to first process"),
     ),
 )
 def test_set_process_lines_with_fixed_and_dynamic_lines(
@@ -144,9 +139,21 @@ def test_set_process_lines_with_fixed_and_dynamic_lines(
     output = ProcessGroupOutput(
         id=1,
         processes=[
-            ProcessOutput(id=1, process=Process(1, "echo first; echo second", lines1)),
-            ProcessOutput(id=2, process=Process(2, "echo first; echo second", lines2)),
-            ProcessOutput(id=3, process=Process(3, "echo first; echo second", lines3)),
+            ProcessOutput(
+                id=1,
+                process=Process(1, "echo first; echo second", lines1),
+                data="first\nsecond\n",
+            ),
+            ProcessOutput(
+                id=2,
+                process=Process(2, "echo first; echo second", lines2),
+                data="first\nsecond\n",
+            ),
+            ProcessOutput(
+                id=3,
+                process=Process(3, "echo first; echo second", lines3),
+                data="first\nsecond\n",
+            ),
         ],
     )