diff --git a/src/prompt_toolkit/layout/containers.py b/src/prompt_toolkit/layout/containers.py index fbc54d8c0..b17b8d202 100644 --- a/src/prompt_toolkit/layout/containers.py +++ b/src/prompt_toolkit/layout/containers.py @@ -448,13 +448,20 @@ def _divide_heights(self, write_position: WritePosition) -> list[int] | None: i = next(child_generator) + # The height handed out so far, carried rather than re-added. + # These loops run once per unit handed out, and `sum(sizes)` + # walked the whole list on every turn to learn what the turn + # before it already knew. + taken = sum(sizes) + # Increase until we meet at least the 'preferred' size. preferred_stop = min(height, sum_dimensions.preferred) preferred_dimensions = [d.preferred for d in dimensions] - while sum(sizes) < preferred_stop: + while taken < preferred_stop: if sizes[i] < preferred_dimensions[i]: sizes[i] += 1 + taken += 1 i = next(child_generator) # Increase until we use all the available space. (or until "max") @@ -462,9 +469,10 @@ def _divide_heights(self, write_position: WritePosition) -> list[int] | None: max_stop = min(height, sum_dimensions.max) max_dimensions = [d.max for d in dimensions] - while sum(sizes) < max_stop: + while taken < max_stop: if sizes[i] < max_dimensions[i]: sizes[i] += 1 + taken += 1 i = next(child_generator) return sizes @@ -643,21 +651,30 @@ def _divide_widths(self, width: int) -> list[int] | None: i = next(child_generator) + # The width handed out so far, carried rather than re-added. + # These loops run about once per column, and `sum(sizes)` walked + # the whole list on every turn to learn what the turn before it + # already knew. + taken = sum(sizes) + # Increase until we meet at least the 'preferred' size. preferred_stop = min(width, sum_dimensions.preferred) + preferred_dimensions = [d.preferred for d in dimensions] - while sum(sizes) < preferred_stop: + while taken < preferred_stop: if sizes[i] < preferred_dimensions[i]: sizes[i] += 1 + taken += 1 i = next(child_generator) # Increase until we use all the available space. max_dimensions = [d.max for d in dimensions] max_stop = min(width, sum_dimensions.max) - while sum(sizes) < max_stop: + while taken < max_stop: if sizes[i] < max_dimensions[i]: sizes[i] += 1 + taken += 1 i = next(child_generator) return sizes diff --git a/src/prompt_toolkit/utils.py b/src/prompt_toolkit/utils.py index 142142a0f..d490e04ab 100644 --- a/src/prompt_toolkit/utils.py +++ b/src/prompt_toolkit/utils.py @@ -249,8 +249,8 @@ def take_using_weights( assert len(items) > 0 # Remove items with zero-weight. - items2 = [] - weights2 = [] + items2: list[_T] = [] + weights2: list[int] = [] for item, w in zip(items, weights): if w > 0: items2.append(item) @@ -264,8 +264,17 @@ def take_using_weights( raise ValueError("Did't got any items with a positive weight.") # - already_taken = [0 for i in items] item_count = len(items) + + # One item takes every turn. `VSplit` and `HSplit` ask for a few + # thousand items of this per render, and a split of one child is the + # common case, so it is worth not building the state below for it. + if item_count == 1: + item = items[0] + while True: + yield item + + already_taken = [0] * item_count max_weight = max(weights) i = 0 @@ -275,9 +284,12 @@ def take_using_weights( while adding: adding = False - for item_i, item, weight in zip(range(item_count), items, weights): - if already_taken[item_i] < i * weight / float(max_weight): - yield item + # Reading the two lists by index, rather than building a + # `zip` object on every pass of a loop that runs once per + # cell handed out. (`range` builds nothing.) + for item_i in range(item_count): + if already_taken[item_i] < i * weights[item_i] / max_weight: + yield items[item_i] already_taken[item_i] += 1 adding = True diff --git a/tests/test_layout.py b/tests/test_layout.py index cbbbcd0bd..bd0a7eae4 100644 --- a/tests/test_layout.py +++ b/tests/test_layout.py @@ -1,10 +1,13 @@ from __future__ import annotations +from random import Random + import pytest from prompt_toolkit.layout import InvalidLayoutError, Layout from prompt_toolkit.layout.containers import HSplit, VSplit, Window from prompt_toolkit.layout.controls import BufferControl +from prompt_toolkit.layout.dimension import Dimension def test_layout_class(): @@ -51,3 +54,46 @@ def test_layout_class(): def test_create_invalid_layout(): with pytest.raises(InvalidLayoutError): Layout(HSplit([])) + + +def test_divide_widths_hands_out_exactly_what_fits(): + """ + `_divide_widths` grows each child one unit at a time and carries the + handed-out total in a local rather than re-adding it. Whatever the + loop does, the answer keeps three promises: too little room gives + back nothing, every child stays inside its own minimum and maximum, + and the sizes add up to everything the split may take. A local that + drifted from the true total in either direction would break the + third promise, so this pins the carried total as well as the loop. + """ + rng = Random(0) + + for _ in range(200): + children = [] + for _ in range(rng.randint(1, 4)): + minimum = rng.randint(0, 5) + preferred = minimum + rng.randint(0, 10) + children.append( + Window( + width=Dimension( + min=minimum, + preferred=preferred, + max=preferred + rng.randint(0, 10), + weight=rng.randint(1, 3), + ) + ) + ) + + width = rng.randint(0, 30) + split = VSplit(children) + actual = [c.preferred_width(width) for c in split._all_children] + sizes = split._divide_widths(width) + + if sum(d.min for d in actual) > width: + assert sizes is None + continue + + assert sizes is not None + assert len(sizes) == len(actual) + assert all(d.min <= s <= d.max for s, d in zip(sizes, actual)) + assert sum(sizes) == min(width, sum(d.max for d in actual)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 9d4c80876..24c19a39c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -76,3 +76,29 @@ def take(generator, count): # All zero-weight items. with pytest.raises(ValueError): take(take_using_weights(["A", "B", "C"], [0, 0, 0]), 70) + + +def test_using_weights_yields_one_exact_order(): + """ + The order, and not only the proportions. + + `VSplit` and `HSplit` hand out one cell at a time in this order, so + two orders with the same proportions still give two layouts: they + differ by which child gets the last cell. + """ + for weights, expected in [ + ([1], "AAAAAAAAAAAA"), + ([7], "AAAAAAAAAAAA"), + ([1, 1], "ABABABABABABABAB"), + ([1, 0], "AAAAAAAAAAAA"), + ([5, 10], "ABBABBABBABBABBABBABBABB"), + ([1, 1, 1], "ABCABCABCABCABCABCABCABC"), + ([3, 1, 2], "ABCACAABCACAABCACAABCACA"), + ([20, 10, 5], "ABCAABAABCAABAABCAABAABCAABAABCAABA"), + ([1, 2, 3, 4], "ABCDCDBCDDABCDCDBCDDABCDCDBCDD"), + ([0, 2, 0, 4], "BDDBDDBDDBDDBDDBDDBDDBDD"), + ]: + items = list("ABCDEFGH"[: len(weights)]) + taken = itertools.islice(take_using_weights(items, weights), len(expected)) + + assert "".join(taken) == expected, weights