xref: /aosp_15_r20/external/pigweed/pw_console/py/pw_console/text_formatting.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2021 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Text formatting functions."""
15
16import copy
17import re
18from typing import Iterable
19
20from prompt_toolkit.formatted_text import StyleAndTextTuples
21from prompt_toolkit.formatted_text.base import OneStyleAndTextTuple
22from prompt_toolkit.utils import get_cwidth
23
24_ANSI_SEQUENCE_REGEX = re.compile(r'\x1b[^m]*m')
25
26
27def strip_ansi(text: str):
28    """Strip out ANSI escape sequences."""
29    return _ANSI_SEQUENCE_REGEX.sub('', text)
30
31
32def split_lines(
33    input_fragments: StyleAndTextTuples,
34) -> list[StyleAndTextTuples]:
35    """Break a flattened list of StyleAndTextTuples into a list of lines.
36
37    Ending line breaks are not preserved."""
38    lines: list[StyleAndTextTuples] = []
39    this_line: StyleAndTextTuples = []
40    for item in input_fragments:
41        if item[1].endswith('\n'):
42            # If there are no elements in this line except for a linebreak add
43            # an empty StyleAndTextTuple so this line isn't an empty list.
44            if len(this_line) == 0 and item[1] == '\n':
45                this_line.append((item[0], item[1][:-1]))
46            lines.append(this_line)
47            this_line = []
48        else:
49            this_line.append(item)
50    return lines
51
52
53def insert_linebreaks(
54    input_fragments: StyleAndTextTuples,
55    max_line_width: int,
56    truncate_long_lines: bool = True,
57) -> tuple[StyleAndTextTuples, int]:
58    """Add line breaks at max_line_width if truncate_long_lines is True.
59
60    Returns input_fragments with each character as it's own formatted text
61    tuple."""
62    fragments: StyleAndTextTuples = []
63    total_width = 0
64    line_width = 0
65    line_height = 0
66    new_break_inserted = False
67
68    for item in input_fragments:
69        # Check for non-printable fragment; doesn't affect the width.
70        if '[ZeroWidthEscape]' in item[0]:
71            fragments.append(item)
72            continue
73
74        new_item_style = item[0]
75
76        # For each character in the fragment
77        for character in item[1]:
78            # Get the width respecting double width characters
79            width = get_cwidth(character)
80            # Increment counters
81            total_width += width
82            line_width += width
83            # Save this character as it's own fragment
84            if line_width <= max_line_width:
85                if not new_break_inserted or character != '\n':
86                    fragments.append((new_item_style, character))
87                    # Was a line break just inserted?
88                    if character == '\n':
89                        # Increase height
90                        line_height += 1
91                new_break_inserted = False
92
93            # Reset width to zero even if we are beyond the max line width.
94            if character == '\n':
95                line_width = 0
96
97            # Are we at the limit for this line?
98            elif line_width == max_line_width:
99                # Insert a new linebreak fragment
100                fragments.append((new_item_style, '\n'))
101                # Increase height
102                line_height += 1
103                # Set a flag for skipping the next character if it is also a
104                # line break.
105                new_break_inserted = True
106
107                if not truncate_long_lines:
108                    # Reset line width to zero
109                    line_width = 0
110
111    # Check if the string ends in a final line break
112    last_fragment_style = fragments[-1][0]
113    last_fragment_text = fragments[-1][1]
114    if not last_fragment_text.endswith('\n'):
115        # Add a line break if none exists
116        fragments.append((last_fragment_style, '\n'))
117        line_height += 1
118
119    return fragments, line_height
120
121
122def join_adjacent_style_tuples(
123    fragments: StyleAndTextTuples,
124) -> StyleAndTextTuples:
125    """Join adjacent FormattedTextTuples if they have the same style."""
126    new_fragments: StyleAndTextTuples = []
127
128    for i, fragment in enumerate(fragments):
129        # Add the first fragment
130        if i == 0:
131            new_fragments.append(fragment)
132            continue
133
134        # Get this style
135        style = fragment[0]
136        # If the previous style matches
137        if style == new_fragments[-1][0]:
138            # Get the previous text
139            new_text = new_fragments[-1][1]
140            # Append this text
141            new_text += fragment[1]
142            # Replace the last fragment
143            new_fragments[-1] = (style, new_text)
144        else:
145            # Styles don't match, just append.
146            new_fragments.append(fragment)
147
148    return new_fragments
149
150
151def fill_character_width(
152    input_fragments: StyleAndTextTuples,
153    fragment_width: int,
154    window_width: int,
155    line_wrapping: bool = False,
156    remaining_width: int = 0,
157    horizontal_scroll_amount: int = 0,
158    add_cursor: bool = False,
159) -> StyleAndTextTuples:
160    """Fill line to the width of the window using spaces."""
161    # Calculate the number of spaces to add at the end.
162    empty_characters = window_width - fragment_width
163    # If wrapping is on, use remaining_width
164    if line_wrapping and (fragment_width > window_width):
165        empty_characters = remaining_width
166
167    # Add additional spaces for horizontal scrolling.
168    empty_characters += horizontal_scroll_amount
169
170    if empty_characters <= 0:
171        # No additional spaces required
172        return input_fragments
173
174    line_fragments = copy.copy(input_fragments)
175
176    single_space = ('', ' ')
177    line_ends_in_a_break = False
178    # Replace the trailing \n with a space
179    if line_fragments[-1][1] == '\n':
180        line_fragments[-1] = single_space
181        empty_characters -= 1
182        line_ends_in_a_break = True
183
184    # Append remaining spaces
185    for _i in range(empty_characters):
186        line_fragments.append(single_space)
187
188    if line_ends_in_a_break:
189        # Restore the \n
190        line_fragments.append(('', '\n'))
191
192    if add_cursor:
193        # Add a cursor to this line by adding SetCursorPosition fragment.
194        line_fragments_remainder = line_fragments
195        line_fragments = [('[SetCursorPosition]', '')]
196        # Use extend to keep types happy.
197        line_fragments.extend(line_fragments_remainder)
198
199    return line_fragments
200
201
202def flatten_formatted_text_tuples(
203    lines: Iterable[StyleAndTextTuples],
204) -> StyleAndTextTuples:
205    """Flatten a list of lines of FormattedTextTuples
206
207    This function will also remove trailing newlines to avoid displaying extra
208    empty lines in prompt_toolkit containers.
209    """
210    fragments: StyleAndTextTuples = []
211
212    # Return empty list if lines is empty.
213    if not lines:
214        return fragments
215
216    for line_fragments in lines:
217        # Append all FormattedText tuples for this line.
218        for fragment in line_fragments:
219            fragments.append(fragment)
220
221    # Strip off any trailing line breaks
222    last_fragment: OneStyleAndTextTuple = fragments[-1]
223    style = last_fragment[0]
224    text = last_fragment[1].rstrip('\n')
225    fragments[-1] = (style, text)
226    return fragments
227
228
229def remove_formatting(formatted_text: StyleAndTextTuples) -> str:
230    """Throw away style info from prompt_toolkit formatted text tuples."""
231    return ''.join([formatted_tuple[1] for formatted_tuple in formatted_text])
232
233
234def get_line_height(text_width, screen_width, prefix_width):
235    """Calculates line height for a string with line wrapping enabled."""
236    if text_width == 0:
237        return 0
238
239    # If text will fit on the screen without wrapping.
240    if text_width <= screen_width:
241        return 1, screen_width - text_width
242
243    # Assume zero width prefix if it's >= width of the screen.
244    if prefix_width >= screen_width:
245        prefix_width = 0
246
247    # Start with height of 1 row.
248    total_height = 1
249
250    # One screen_width of characters (with no prefix) is displayed first.
251    remaining_width = text_width - screen_width
252
253    # While we have caracters remaining to be displayed
254    while remaining_width > 0:
255        # Add the new indentation prefix
256        remaining_width += prefix_width
257        # Display this line
258        remaining_width -= screen_width
259        # Add a line break
260        total_height += 1
261
262    # Remaining characters is what's left below zero.
263    return (total_height, abs(remaining_width))
264