From d29f5d5464c7fd0451e9ccaa4427a1e9b3f13142 Mon Sep 17 00:00:00 2001 From: Danthewaann Date: Fri, 21 Feb 2025 18:03:51 +0000 Subject: [PATCH] fix: allocate more lines to processes with more output --- src/pyallel/printer.py | 74 +++++++++++++++++--------------- src/pyallel/process.py | 2 + tests/test_printer.py | 95 +++++++++++++++++++++++------------------- 3 files changed, 93 insertions(+), 78 deletions(-) diff --git a/src/pyallel/printer.py b/src/pyallel/printer.py index 08e2c01..57a829c 100644 --- a/src/pyallel/printer.py +++ b/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: diff --git a/src/pyallel/process.py b/src/pyallel/process.py index 58379cd..9abf407 100644 --- a/src/pyallel/process.py +++ b/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: diff --git a/tests/test_printer.py b/tests/test_printer.py index 0a22e28..984a49b 100644 --- a/tests/test_printer.py +++ b/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", + ), ], )