1# coding=utf-8
2"""
3cmd2 table creation API
4This API is built upon two core classes: Column and TableCreator
5The general use case is to inherit from TableCreator to create a table class with custom formatting options.
6There are already implemented and ready-to-use examples of this below TableCreator's code.
7"""
8import copy
9import io
10from collections import (
11    deque,
12)
13from enum import (
14    Enum,
15)
16from typing import (
17    Any,
18    Deque,
19    List,
20    Optional,
21    Sequence,
22    Tuple,
23    Union,
24)
25
26from wcwidth import (  # type: ignore[import]
27    wcwidth,
28)
29
30from . import (
31    ansi,
32    constants,
33    utils,
34)
35
36# Constants
37EMPTY = ''
38SPACE = ' '
39
40
41class HorizontalAlignment(Enum):
42    """Horizontal alignment of text in a cell"""
43
44    LEFT = 1
45    CENTER = 2
46    RIGHT = 3
47
48
49class VerticalAlignment(Enum):
50    """Vertical alignment of text in a cell"""
51
52    TOP = 1
53    MIDDLE = 2
54    BOTTOM = 3
55
56
57class Column:
58    """Table column configuration"""
59
60    def __init__(
61        self,
62        header: str,
63        *,
64        width: Optional[int] = None,
65        header_horiz_align: HorizontalAlignment = HorizontalAlignment.LEFT,
66        header_vert_align: VerticalAlignment = VerticalAlignment.BOTTOM,
67        style_header_text: bool = True,
68        data_horiz_align: HorizontalAlignment = HorizontalAlignment.LEFT,
69        data_vert_align: VerticalAlignment = VerticalAlignment.TOP,
70        style_data_text: bool = True,
71        max_data_lines: Union[int, float] = constants.INFINITY,
72    ) -> None:
73        """
74        Column initializer
75
76        :param header: label for column header
77        :param width: display width of column. This does not account for any borders or padding which
78                      may be added (e.g pre_line, inter_cell, and post_line). Header and data text wrap within
79                      this width using word-based wrapping (defaults to actual width of header or 1 if header is blank)
80        :param header_horiz_align: horizontal alignment of header cells (defaults to left)
81        :param header_vert_align: vertical alignment of header cells (defaults to bottom)
82        :param style_header_text: if True, then the table is allowed to apply styles to the header text, which may
83                                  conflict with any styles the header already has. If False, the header is printed as is.
84                                  Table classes which apply style to headers must account for the value of this flag.
85                                  (defaults to True)
86        :param data_horiz_align: horizontal alignment of data cells (defaults to left)
87        :param data_vert_align: vertical alignment of data cells (defaults to top)
88        :param style_data_text: if True, then the table is allowed to apply styles to the data text, which may
89                                conflict with any styles the data already has. If False, the data is printed as is.
90                                Table classes which apply style to data must account for the value of this flag.
91                                (defaults to True)
92        :param max_data_lines: maximum lines allowed in a data cell. If line count exceeds this, then the final
93                               line displayed will be truncated with an ellipsis. (defaults to INFINITY)
94        :raises: ValueError if width is less than 1
95        :raises: ValueError if max_data_lines is less than 1
96        """
97        self.header = header
98
99        if width is not None and width < 1:
100            raise ValueError("Column width cannot be less than 1")
101        else:
102            self.width: int = width if width is not None else -1
103
104        self.header_horiz_align = header_horiz_align
105        self.header_vert_align = header_vert_align
106        self.style_header_text = style_header_text
107
108        self.data_horiz_align = data_horiz_align
109        self.data_vert_align = data_vert_align
110        self.style_data_text = style_data_text
111
112        if max_data_lines < 1:
113            raise ValueError("Max data lines cannot be less than 1")
114
115        self.max_data_lines = max_data_lines
116
117
118class TableCreator:
119    """
120    Base table creation class. This class handles ANSI style sequences and characters with display widths greater than 1
121    when performing width calculations. It was designed with the ability to build tables one row at a time. This helps
122    when you have large data sets that you don't want to hold in memory or when you receive portions of the data set
123    incrementally.
124
125    TableCreator has one public method: generate_row()
126
127    This function and the Column class provide all features needed to build tables with headers, borders, colors,
128    horizontal and vertical alignment, and wrapped text. However, it's generally easier to inherit from this class and
129    implement a more granular API rather than use TableCreator directly. There are ready-to-use examples of this
130    defined after this class.
131    """
132
133    def __init__(self, cols: Sequence[Column], *, tab_width: int = 4) -> None:
134        """
135        TableCreator initializer
136
137        :param cols: column definitions for this table
138        :param tab_width: all tabs will be replaced with this many spaces. If a row's fill_char is a tab,
139                          then it will be converted to one space.
140        :raises: ValueError if tab_width is less than 1
141        """
142        if tab_width < 1:
143            raise ValueError("Tab width cannot be less than 1")
144
145        self.cols = copy.copy(cols)
146        self.tab_width = tab_width
147
148        for col in self.cols:
149            # Replace tabs before calculating width of header strings
150            col.header = col.header.replace('\t', SPACE * self.tab_width)
151
152            # For headers with the width not yet set, use the width of the
153            # widest line in the header or 1 if the header has no width
154            if col.width <= 0:
155                col.width = max(1, ansi.widest_line(col.header))
156
157    @staticmethod
158    def _wrap_long_word(word: str, max_width: int, max_lines: Union[int, float], is_last_word: bool) -> Tuple[str, int, int]:
159        """
160        Used by _wrap_text() to wrap a long word over multiple lines
161
162        :param word: word being wrapped
163        :param max_width: maximum display width of a line
164        :param max_lines: maximum lines to wrap before ending the last line displayed with an ellipsis
165        :param is_last_word: True if this is the last word of the total text being wrapped
166        :return: Tuple(wrapped text, lines used, display width of last line)
167        """
168        styles = utils.get_styles_in_text(word)
169        wrapped_buf = io.StringIO()
170
171        # How many lines we've used
172        total_lines = 1
173
174        # Display width of the current line we are building
175        cur_line_width = 0
176
177        char_index = 0
178        while char_index < len(word):
179            # We've reached the last line. Let truncate_line do the rest.
180            if total_lines == max_lines:
181                # If this isn't the last word, but it's gonna fill the final line, then force truncate_line
182                # to place an ellipsis at the end of it by making the word too wide.
183                remaining_word = word[char_index:]
184                if not is_last_word and ansi.style_aware_wcswidth(remaining_word) == max_width:
185                    remaining_word += "EXTRA"
186
187                truncated_line = utils.truncate_line(remaining_word, max_width)
188                cur_line_width = ansi.style_aware_wcswidth(truncated_line)
189                wrapped_buf.write(truncated_line)
190                break
191
192            # Check if we're at a style sequence. These don't count toward display width.
193            if char_index in styles:
194                wrapped_buf.write(styles[char_index])
195                char_index += len(styles[char_index])
196                continue
197
198            cur_char = word[char_index]
199            cur_char_width = wcwidth(cur_char)
200
201            if cur_char_width > max_width:
202                # We have a case where the character is wider than max_width. This can happen if max_width
203                # is 1 and the text contains wide characters (e.g. East Asian). Replace it with an ellipsis.
204                cur_char = constants.HORIZONTAL_ELLIPSIS
205                cur_char_width = wcwidth(cur_char)
206
207            if cur_line_width + cur_char_width > max_width:
208                # Adding this char will exceed the max_width. Start a new line.
209                wrapped_buf.write('\n')
210                total_lines += 1
211                cur_line_width = 0
212                continue
213
214            # Add this character and move to the next one
215            cur_line_width += cur_char_width
216            wrapped_buf.write(cur_char)
217            char_index += 1
218
219        return wrapped_buf.getvalue(), total_lines, cur_line_width
220
221    @staticmethod
222    def _wrap_text(text: str, max_width: int, max_lines: Union[int, float]) -> str:
223        """
224        Wrap text into lines with a display width no longer than max_width. This function breaks words on whitespace
225        boundaries. If a word is longer than the space remaining on a line, then it will start on a new line.
226        ANSI escape sequences do not count toward the width of a line.
227
228        :param text: text to be wrapped
229        :param max_width: maximum display width of a line
230        :param max_lines: maximum lines to wrap before ending the last line displayed with an ellipsis
231        :return: wrapped text
232        """
233
234        # MyPy Issue #7057 documents regression requiring nonlocals to be defined earlier
235        cur_line_width = 0
236        total_lines = 0
237
238        def add_word(word_to_add: str, is_last_word: bool) -> None:
239            """
240            Called from loop to add a word to the wrapped text
241
242            :param word_to_add: the word being added
243            :param is_last_word: True if this is the last word of the total text being wrapped
244            """
245            nonlocal cur_line_width
246            nonlocal total_lines
247
248            # No more space to add word
249            if total_lines == max_lines and cur_line_width == max_width:
250                return
251
252            word_width = ansi.style_aware_wcswidth(word_to_add)
253
254            # If the word is wider than max width of a line, attempt to start it on its own line and wrap it
255            if word_width > max_width:
256                room_to_add = True
257
258                if cur_line_width > 0:
259                    # The current line already has text, check if there is room to create a new line
260                    if total_lines < max_lines:
261                        wrapped_buf.write('\n')
262                        total_lines += 1
263                    else:
264                        # We will truncate this word on the remaining line
265                        room_to_add = False
266
267                if room_to_add:
268                    wrapped_word, lines_used, cur_line_width = TableCreator._wrap_long_word(
269                        word_to_add, max_width, max_lines - total_lines + 1, is_last_word
270                    )
271                    # Write the word to the buffer
272                    wrapped_buf.write(wrapped_word)
273                    total_lines += lines_used - 1
274                    return
275
276            # We aren't going to wrap the word across multiple lines
277            remaining_width = max_width - cur_line_width
278
279            # Check if we need to start a new line
280            if word_width > remaining_width and total_lines < max_lines:
281                # Save the last character in wrapped_buf, which can't be empty at this point.
282                seek_pos = wrapped_buf.tell() - 1
283                wrapped_buf.seek(seek_pos)
284                last_char = wrapped_buf.read()
285
286                wrapped_buf.write('\n')
287                total_lines += 1
288                cur_line_width = 0
289                remaining_width = max_width
290
291                # Only when a space is following a space do we want to start the next line with it.
292                if word_to_add == SPACE and last_char != SPACE:
293                    return
294
295            # Check if we've hit the last line we're allowed to create
296            if total_lines == max_lines:
297                # If this word won't fit, truncate it
298                if word_width > remaining_width:
299                    word_to_add = utils.truncate_line(word_to_add, remaining_width)
300                    word_width = remaining_width
301
302                # If this isn't the last word, but it's gonna fill the final line, then force truncate_line
303                # to place an ellipsis at the end of it by making the word too wide.
304                elif not is_last_word and word_width == remaining_width:
305                    word_to_add = utils.truncate_line(word_to_add + "EXTRA", remaining_width)
306
307            cur_line_width += word_width
308            wrapped_buf.write(word_to_add)
309
310        ############################################################################################################
311        # _wrap_text() main code
312        ############################################################################################################
313        # Buffer of the wrapped text
314        wrapped_buf = io.StringIO()
315
316        # How many lines we've used
317        total_lines = 0
318
319        # Respect the existing line breaks
320        data_str_lines = text.splitlines()
321        for data_line_index, data_line in enumerate(data_str_lines):
322            total_lines += 1
323
324            if data_line_index > 0:
325                wrapped_buf.write('\n')
326
327            # If the last line is empty, then add a newline and stop
328            if data_line_index == len(data_str_lines) - 1 and not data_line:
329                wrapped_buf.write('\n')
330                break
331
332            # Locate the styles in this line
333            styles = utils.get_styles_in_text(data_line)
334
335            # Display width of the current line we are building
336            cur_line_width = 0
337
338            # Current word being built
339            cur_word_buf = io.StringIO()
340
341            char_index = 0
342            while char_index < len(data_line):
343                if total_lines == max_lines and cur_line_width == max_width:
344                    break
345
346                # Check if we're at a style sequence. These don't count toward display width.
347                if char_index in styles:
348                    cur_word_buf.write(styles[char_index])
349                    char_index += len(styles[char_index])
350                    continue
351
352                cur_char = data_line[char_index]
353                if cur_char == SPACE:
354                    # If we've reached the end of a word, then add the word to the wrapped text
355                    if cur_word_buf.tell() > 0:
356                        # is_last_word is False since there is a space after the word
357                        add_word(cur_word_buf.getvalue(), is_last_word=False)
358                        cur_word_buf = io.StringIO()
359
360                    # Add the space to the wrapped text
361                    last_word = data_line_index == len(data_str_lines) - 1 and char_index == len(data_line) - 1
362                    add_word(cur_char, last_word)
363                else:
364                    # Add this character to the word buffer
365                    cur_word_buf.write(cur_char)
366
367                char_index += 1
368
369            # Add the final word of this line if it's been started
370            if cur_word_buf.tell() > 0:
371                last_word = data_line_index == len(data_str_lines) - 1 and char_index == len(data_line)
372                add_word(cur_word_buf.getvalue(), last_word)
373
374            # Stop line loop if we've written to max_lines
375            if total_lines == max_lines:
376                # If this isn't the last data line and there is space
377                # left on the final wrapped line, then add an ellipsis
378                if data_line_index < len(data_str_lines) - 1 and cur_line_width < max_width:
379                    wrapped_buf.write(constants.HORIZONTAL_ELLIPSIS)
380                break
381
382        return wrapped_buf.getvalue()
383
384    def _generate_cell_lines(self, cell_data: Any, is_header: bool, col: Column, fill_char: str) -> Tuple[Deque[str], int]:
385        """
386        Generate the lines of a table cell
387
388        :param cell_data: data to be included in cell
389        :param is_header: True if writing a header cell, otherwise writing a data cell. This determines whether to
390                          use header or data alignment settings as well as maximum lines to wrap.
391        :param col: Column definition for this cell
392        :param fill_char: character that fills remaining space in a cell. If your text has a background color,
393                          then give fill_char the same background color. (Cannot be a line breaking character)
394        :return: Tuple of cell lines deque and the display width of the cell
395        """
396        # Convert data to string and replace tabs with spaces
397        data_str = str(cell_data).replace('\t', SPACE * self.tab_width)
398
399        # Wrap text in this cell
400        max_lines = constants.INFINITY if is_header else col.max_data_lines
401        wrapped_text = self._wrap_text(data_str, col.width, max_lines)
402
403        # Align the text horizontally
404        horiz_alignment = col.header_horiz_align if is_header else col.data_horiz_align
405        if horiz_alignment == HorizontalAlignment.LEFT:
406            text_alignment = utils.TextAlignment.LEFT
407        elif horiz_alignment == HorizontalAlignment.CENTER:
408            text_alignment = utils.TextAlignment.CENTER
409        else:
410            text_alignment = utils.TextAlignment.RIGHT
411
412        aligned_text = utils.align_text(wrapped_text, fill_char=fill_char, width=col.width, alignment=text_alignment)
413
414        lines = deque(aligned_text.splitlines())
415        cell_width = ansi.widest_line(aligned_text)
416        return lines, cell_width
417
418    def generate_row(
419        self,
420        row_data: Sequence[Any],
421        is_header: bool,
422        *,
423        fill_char: str = SPACE,
424        pre_line: str = EMPTY,
425        inter_cell: str = (2 * SPACE),
426        post_line: str = EMPTY,
427    ) -> str:
428        """
429        Generate a header or data table row
430
431        :param row_data: data with an entry for each column in the row
432        :param is_header: True if writing a header cell, otherwise writing a data cell. This determines whether to
433                          use header or data alignment settings as well as maximum lines to wrap.
434        :param fill_char: character that fills remaining space in a cell. Defaults to space. If this is a tab,
435                          then it will be converted to one space. (Cannot be a line breaking character)
436        :param pre_line: string to print before each line of a row. This can be used for a left row border and
437                         padding before the first cell's text. (Defaults to blank)
438        :param inter_cell: string to print where two cells meet. This can be used for a border between cells and padding
439                           between it and the 2 cells' text. (Defaults to 2 spaces)
440        :param post_line: string to print after each line of a row. This can be used for padding after
441                          the last cell's text and a right row border. (Defaults to blank)
442        :return: row string
443        :raises: ValueError if data isn't the same length as self.cols
444        :raises: TypeError if fill_char is more than one character (not including ANSI style sequences)
445        :raises: ValueError if fill_char, pre_line, inter_cell, or post_line contains an unprintable
446                 character like a newline
447        """
448
449        class Cell:
450            """Inner class which represents a table cell"""
451
452            def __init__(self) -> None:
453                # Data in this cell split into individual lines
454                self.lines: Deque[str] = deque()
455
456                # Display width of this cell
457                self.width = 0
458
459        if len(row_data) != len(self.cols):
460            raise ValueError("Length of row_data must match length of cols")
461
462        # Replace tabs (tabs in data strings will be handled in _generate_cell_lines())
463        fill_char = fill_char.replace('\t', SPACE)
464        pre_line = pre_line.replace('\t', SPACE * self.tab_width)
465        inter_cell = inter_cell.replace('\t', SPACE * self.tab_width)
466        post_line = post_line.replace('\t', SPACE * self.tab_width)
467
468        # Validate fill_char character count
469        if len(ansi.strip_style(fill_char)) != 1:
470            raise TypeError("Fill character must be exactly one character long")
471
472        # Look for unprintable characters
473        validation_dict = {'fill_char': fill_char, 'pre_line': pre_line, 'inter_cell': inter_cell, 'post_line': post_line}
474        for key, val in validation_dict.items():
475            if ansi.style_aware_wcswidth(val) == -1:
476                raise ValueError(f"{key} contains an unprintable character")
477
478        # Number of lines this row uses
479        total_lines = 0
480
481        # Generate the cells for this row
482        cells = list()
483
484        for col_index, col in enumerate(self.cols):
485            cell = Cell()
486            cell.lines, cell.width = self._generate_cell_lines(row_data[col_index], is_header, col, fill_char)
487            cells.append(cell)
488            total_lines = max(len(cell.lines), total_lines)
489
490        row_buf = io.StringIO()
491
492        # Vertically align each cell
493        for cell_index, cell in enumerate(cells):
494            col = self.cols[cell_index]
495            vert_align = col.header_vert_align if is_header else col.data_vert_align
496
497            # Check if this cell need vertical filler
498            line_diff = total_lines - len(cell.lines)
499            if line_diff == 0:
500                continue
501
502            # Add vertical filler lines
503            padding_line = utils.align_left(EMPTY, fill_char=fill_char, width=cell.width)
504            if vert_align == VerticalAlignment.TOP:
505                to_top = 0
506                to_bottom = line_diff
507            elif vert_align == VerticalAlignment.MIDDLE:
508                to_top = line_diff // 2
509                to_bottom = line_diff - to_top
510            else:
511                to_top = line_diff
512                to_bottom = 0
513
514            for i in range(to_top):
515                cell.lines.appendleft(padding_line)
516            for i in range(to_bottom):
517                cell.lines.append(padding_line)
518
519        # Build this row one line at a time
520        for line_index in range(total_lines):
521            for cell_index, cell in enumerate(cells):
522                if cell_index == 0:
523                    row_buf.write(pre_line)
524
525                row_buf.write(cell.lines[line_index])
526
527                if cell_index < len(self.cols) - 1:
528                    row_buf.write(inter_cell)
529                if cell_index == len(self.cols) - 1:
530                    row_buf.write(post_line)
531
532            # Add a newline if this is not the last line
533            if line_index < total_lines - 1:
534                row_buf.write('\n')
535
536        return row_buf.getvalue()
537
538
539############################################################################################################
540# The following are implementations of TableCreator which demonstrate how to make various types
541# of tables. They can be used as-is or serve as inspiration for other custom table classes.
542############################################################################################################
543class SimpleTable(TableCreator):
544    """
545    Implementation of TableCreator which generates a borderless table with an optional divider row after the header.
546    This class can be used to create the whole table at once or one row at a time.
547    """
548
549    def __init__(
550        self,
551        cols: Sequence[Column],
552        *,
553        column_spacing: int = 2,
554        tab_width: int = 4,
555        divider_char: Optional[str] = '-',
556        header_bg: Optional[ansi.BgColor] = None,
557        data_bg: Optional[ansi.BgColor] = None,
558    ) -> None:
559        """
560        SimpleTable initializer
561
562        :param cols: column definitions for this table
563        :param column_spacing: how many spaces to place between columns. Defaults to 2.
564        :param tab_width: all tabs will be replaced with this many spaces. If a row's fill_char is a tab,
565                          then it will be converted to one space.
566        :param divider_char: optional character used to build the header divider row. Set this to blank or None if you don't
567                             want a divider row. Defaults to dash. (Cannot be a line breaking character)
568        :param header_bg: optional background color for header cells (defaults to None)
569        :param data_bg: optional background color for data cells (defaults to None)
570        :raises: ValueError if tab_width is less than 1
571        :raises: ValueError if column_spacing is less than 0
572        :raises: TypeError if divider_char is longer than one character
573        :raises: ValueError if divider_char is an unprintable character
574        """
575        super().__init__(cols, tab_width=tab_width)
576
577        if column_spacing < 0:
578            raise ValueError("Column spacing cannot be less than 0")
579
580        self.column_spacing = column_spacing
581
582        if divider_char == '':
583            divider_char = None
584
585        if divider_char is not None:
586            if len(ansi.strip_style(divider_char)) != 1:
587                raise TypeError("Divider character must be exactly one character long")
588
589            divider_char_width = ansi.style_aware_wcswidth(divider_char)
590            if divider_char_width == -1:
591                raise ValueError("Divider character is an unprintable character")
592
593        self.divider_char = divider_char
594        self.header_bg = header_bg
595        self.data_bg = data_bg
596
597    def apply_header_bg(self, value: Any) -> str:
598        """
599        If defined, apply the header background color to header text
600        :param value: object whose text is to be colored
601        :return: formatted text
602        """
603        if self.header_bg is None:
604            return str(value)
605        return ansi.style(value, bg=self.header_bg)
606
607    def apply_data_bg(self, value: Any) -> str:
608        """
609        If defined, apply the data background color to data text
610        :param value: object whose text is to be colored
611        :return: formatted data string
612        """
613        if self.data_bg is None:
614            return str(value)
615        return ansi.style(value, bg=self.data_bg)
616
617    @classmethod
618    def base_width(cls, num_cols: int, *, column_spacing: int = 2) -> int:
619        """
620        Utility method to calculate the display width required for a table before data is added to it.
621        This is useful when determining how wide to make your columns to have a table be a specific width.
622
623        :param num_cols: how many columns the table will have
624        :param column_spacing: how many spaces to place between columns. Defaults to 2.
625        :return: base width
626        :raises: ValueError if column_spacing is less than 0
627        :raises: ValueError if num_cols is less than 1
628        """
629        if num_cols < 1:
630            raise ValueError("Column count cannot be less than 1")
631
632        data_str = SPACE
633        data_width = ansi.style_aware_wcswidth(data_str) * num_cols
634
635        tbl = cls([Column(data_str)] * num_cols, column_spacing=column_spacing)
636        data_row = tbl.generate_data_row([data_str] * num_cols)
637
638        return ansi.style_aware_wcswidth(data_row) - data_width
639
640    def total_width(self) -> int:
641        """Calculate the total display width of this table"""
642        base_width = self.base_width(len(self.cols), column_spacing=self.column_spacing)
643        data_width = sum(col.width for col in self.cols)
644        return base_width + data_width
645
646    def generate_header(self) -> str:
647        """Generate table header with an optional divider row"""
648        header_buf = io.StringIO()
649
650        fill_char = self.apply_header_bg(SPACE)
651        inter_cell = self.apply_header_bg(self.column_spacing * SPACE)
652
653        # Apply background color to header text in Columns which allow it
654        to_display: List[Any] = []
655        for col in self.cols:
656            if col.style_header_text:
657                to_display.append(self.apply_header_bg(col.header))
658            else:
659                to_display.append(col.header)
660
661        # Create the header labels
662        header_labels = self.generate_row(to_display, is_header=True, fill_char=fill_char, inter_cell=inter_cell)
663        header_buf.write(header_labels)
664
665        # Add the divider if necessary
666        divider = self.generate_divider()
667        if divider:
668            header_buf.write('\n' + divider)
669
670        return header_buf.getvalue()
671
672    def generate_divider(self) -> str:
673        """Generate divider row"""
674        if self.divider_char is None:
675            return ''
676
677        return utils.align_left('', fill_char=self.divider_char, width=self.total_width())
678
679    def generate_data_row(self, row_data: Sequence[Any]) -> str:
680        """
681        Generate a data row
682
683        :param row_data: data with an entry for each column in the row
684        :return: data row string
685        """
686        fill_char = self.apply_data_bg(SPACE)
687        inter_cell = self.apply_data_bg(self.column_spacing * SPACE)
688
689        # Apply background color to data text in Columns which allow it
690        to_display: List[Any] = []
691        for index, col in enumerate(self.cols):
692            if col.style_data_text:
693                to_display.append(self.apply_data_bg(row_data[index]))
694            else:
695                to_display.append(row_data[index])
696
697        return self.generate_row(to_display, is_header=False, fill_char=fill_char, inter_cell=inter_cell)
698
699    def generate_table(self, table_data: Sequence[Sequence[Any]], *, include_header: bool = True, row_spacing: int = 1) -> str:
700        """
701        Generate a table from a data set
702
703        :param table_data: Data with an entry for each data row of the table. Each entry should have data for
704                           each column in the row.
705        :param include_header: If True, then a header will be included at top of table. (Defaults to True)
706        :param row_spacing: A number 0 or greater specifying how many blank lines to place between
707                            each row (Defaults to 1)
708        :raises: ValueError if row_spacing is less than 0
709        """
710        if row_spacing < 0:
711            raise ValueError("Row spacing cannot be less than 0")
712
713        table_buf = io.StringIO()
714
715        if include_header:
716            header = self.generate_header()
717            table_buf.write(header)
718            if len(table_data) > 0:
719                table_buf.write('\n')
720
721        row_divider = utils.align_left('', fill_char=self.apply_data_bg(SPACE), width=self.total_width()) + '\n'
722
723        for index, row_data in enumerate(table_data):
724            if index > 0 and row_spacing > 0:
725                table_buf.write(row_spacing * row_divider)
726
727            row = self.generate_data_row(row_data)
728            table_buf.write(row)
729            if index < len(table_data) - 1:
730                table_buf.write('\n')
731
732        return table_buf.getvalue()
733
734
735class BorderedTable(TableCreator):
736    """
737    Implementation of TableCreator which generates a table with borders around the table and between rows. Borders
738    between columns can also be toggled. This class can be used to create the whole table at once or one row at a time.
739    """
740
741    def __init__(
742        self,
743        cols: Sequence[Column],
744        *,
745        tab_width: int = 4,
746        column_borders: bool = True,
747        padding: int = 1,
748        border_fg: Optional[ansi.FgColor] = None,
749        border_bg: Optional[ansi.BgColor] = None,
750        header_bg: Optional[ansi.BgColor] = None,
751        data_bg: Optional[ansi.BgColor] = None,
752    ) -> None:
753        """
754        BorderedTable initializer
755
756        :param cols: column definitions for this table
757        :param tab_width: all tabs will be replaced with this many spaces. If a row's fill_char is a tab,
758                          then it will be converted to one space.
759        :param column_borders: if True, borders between columns will be included. This gives the table a grid-like
760                               appearance. Turning off column borders results in a unified appearance between
761                               a row's cells. (Defaults to True)
762        :param padding: number of spaces between text and left/right borders of cell
763        :param border_fg: optional foreground color for borders (defaults to None)
764        :param border_bg: optional background color for borders (defaults to None)
765        :param header_bg: optional background color for header cells (defaults to None)
766        :param data_bg: optional background color for data cells (defaults to None)
767        :raises: ValueError if tab_width is less than 1
768        :raises: ValueError if padding is less than 0
769        """
770        super().__init__(cols, tab_width=tab_width)
771        self.empty_data = [EMPTY] * len(self.cols)
772        self.column_borders = column_borders
773
774        if padding < 0:
775            raise ValueError("Padding cannot be less than 0")
776        self.padding = padding
777
778        self.border_fg = border_fg
779        self.border_bg = border_bg
780        self.header_bg = header_bg
781        self.data_bg = data_bg
782
783    def apply_border_color(self, value: Any) -> str:
784        """
785        If defined, apply the border foreground and background colors
786        :param value: object whose text is to be colored
787        :return: formatted text
788        """
789        if self.border_fg is None and self.border_bg is None:
790            return str(value)
791        return ansi.style(value, fg=self.border_fg, bg=self.border_bg)
792
793    def apply_header_bg(self, value: Any) -> str:
794        """
795        If defined, apply the header background color to header text
796        :param value: object whose text is to be colored
797        :return: formatted text
798        """
799        if self.header_bg is None:
800            return str(value)
801        return ansi.style(value, bg=self.header_bg)
802
803    def apply_data_bg(self, value: Any) -> str:
804        """
805        If defined, apply the data background color to data text
806        :param value: object whose text is to be colored
807        :return: formatted data string
808        """
809        if self.data_bg is None:
810            return str(value)
811        return ansi.style(value, bg=self.data_bg)
812
813    @classmethod
814    def base_width(cls, num_cols: int, *, column_borders: bool = True, padding: int = 1) -> int:
815        """
816        Utility method to calculate the display width required for a table before data is added to it.
817        This is useful when determining how wide to make your columns to have a table be a specific width.
818
819        :param num_cols: how many columns the table will have
820        :param column_borders: if True, borders between columns will be included in the calculation (Defaults to True)
821        :param padding: number of spaces between text and left/right borders of cell
822        :return: base width
823        :raises: ValueError if num_cols is less than 1
824        """
825        if num_cols < 1:
826            raise ValueError("Column count cannot be less than 1")
827
828        data_str = SPACE
829        data_width = ansi.style_aware_wcswidth(data_str) * num_cols
830
831        tbl = cls([Column(data_str)] * num_cols, column_borders=column_borders, padding=padding)
832        data_row = tbl.generate_data_row([data_str] * num_cols)
833
834        return ansi.style_aware_wcswidth(data_row) - data_width
835
836    def total_width(self) -> int:
837        """Calculate the total display width of this table"""
838        base_width = self.base_width(len(self.cols), column_borders=self.column_borders, padding=self.padding)
839        data_width = sum(col.width for col in self.cols)
840        return base_width + data_width
841
842    def generate_table_top_border(self) -> str:
843        """Generate a border which appears at the top of the header and data section"""
844        fill_char = '═'
845
846        pre_line = '╔' + self.padding * '═'
847
848        inter_cell = self.padding * '═'
849        if self.column_borders:
850            inter_cell += "╤"
851        inter_cell += self.padding * '═'
852
853        post_line = self.padding * '═' + '╗'
854
855        return self.generate_row(
856            self.empty_data,
857            is_header=False,
858            fill_char=self.apply_border_color(fill_char),
859            pre_line=self.apply_border_color(pre_line),
860            inter_cell=self.apply_border_color(inter_cell),
861            post_line=self.apply_border_color(post_line),
862        )
863
864    def generate_header_bottom_border(self) -> str:
865        """Generate a border which appears at the bottom of the header"""
866        fill_char = '═'
867
868        pre_line = '╠' + self.padding * '═'
869
870        inter_cell = self.padding * '═'
871        if self.column_borders:
872            inter_cell += '╪'
873        inter_cell += self.padding * '═'
874
875        post_line = self.padding * '═' + '╣'
876
877        return self.generate_row(
878            self.empty_data,
879            is_header=False,
880            fill_char=self.apply_border_color(fill_char),
881            pre_line=self.apply_border_color(pre_line),
882            inter_cell=self.apply_border_color(inter_cell),
883            post_line=self.apply_border_color(post_line),
884        )
885
886    def generate_row_bottom_border(self) -> str:
887        """Generate a border which appears at the bottom of rows"""
888        fill_char = '─'
889
890        pre_line = '╟' + self.padding * '─'
891
892        inter_cell = self.padding * '─'
893        if self.column_borders:
894            inter_cell += '┼'
895        inter_cell += self.padding * '─'
896        inter_cell = inter_cell
897
898        post_line = self.padding * '─' + '╢'
899
900        return self.generate_row(
901            self.empty_data,
902            is_header=False,
903            fill_char=self.apply_border_color(fill_char),
904            pre_line=self.apply_border_color(pre_line),
905            inter_cell=self.apply_border_color(inter_cell),
906            post_line=self.apply_border_color(post_line),
907        )
908
909    def generate_table_bottom_border(self) -> str:
910        """Generate a border which appears at the bottom of the table"""
911        fill_char = '═'
912
913        pre_line = '╚' + self.padding * '═'
914
915        inter_cell = self.padding * '═'
916        if self.column_borders:
917            inter_cell += '╧'
918        inter_cell += self.padding * '═'
919
920        post_line = self.padding * '═' + '╝'
921
922        return self.generate_row(
923            self.empty_data,
924            is_header=False,
925            fill_char=self.apply_border_color(fill_char),
926            pre_line=self.apply_border_color(pre_line),
927            inter_cell=self.apply_border_color(inter_cell),
928            post_line=self.apply_border_color(post_line),
929        )
930
931    def generate_header(self) -> str:
932        """Generate table header"""
933        fill_char = self.apply_header_bg(SPACE)
934
935        pre_line = self.apply_border_color('║') + self.apply_header_bg(self.padding * SPACE)
936
937        inter_cell = self.apply_header_bg(self.padding * SPACE)
938        if self.column_borders:
939            inter_cell += self.apply_border_color('│')
940        inter_cell += self.apply_header_bg(self.padding * SPACE)
941
942        post_line = self.apply_header_bg(self.padding * SPACE) + self.apply_border_color('║')
943
944        # Apply background color to header text in Columns which allow it
945        to_display: List[Any] = []
946        for col in self.cols:
947            if col.style_header_text:
948                to_display.append(self.apply_header_bg(col.header))
949            else:
950                to_display.append(col.header)
951
952        # Create the bordered header
953        header_buf = io.StringIO()
954        header_buf.write(self.generate_table_top_border())
955        header_buf.write('\n')
956        header_buf.write(
957            self.generate_row(
958                to_display, is_header=True, fill_char=fill_char, pre_line=pre_line, inter_cell=inter_cell, post_line=post_line
959            )
960        )
961        header_buf.write('\n')
962        header_buf.write(self.generate_header_bottom_border())
963
964        return header_buf.getvalue()
965
966    def generate_data_row(self, row_data: Sequence[Any]) -> str:
967        """
968        Generate a data row
969
970        :param row_data: data with an entry for each column in the row
971        :return: data row string
972        """
973        fill_char = self.apply_data_bg(SPACE)
974
975        pre_line = self.apply_border_color('║') + self.apply_data_bg(self.padding * SPACE)
976
977        inter_cell = self.apply_data_bg(self.padding * SPACE)
978        if self.column_borders:
979            inter_cell += self.apply_border_color('│')
980        inter_cell += self.apply_data_bg(self.padding * SPACE)
981
982        post_line = self.apply_data_bg(self.padding * SPACE) + self.apply_border_color('║')
983
984        # Apply background color to data text in Columns which allow it
985        to_display: List[Any] = []
986        for index, col in enumerate(self.cols):
987            if col.style_data_text:
988                to_display.append(self.apply_data_bg(row_data[index]))
989            else:
990                to_display.append(row_data[index])
991
992        return self.generate_row(
993            to_display, is_header=False, fill_char=fill_char, pre_line=pre_line, inter_cell=inter_cell, post_line=post_line
994        )
995
996    def generate_table(self, table_data: Sequence[Sequence[Any]], *, include_header: bool = True) -> str:
997        """
998        Generate a table from a data set
999
1000        :param table_data: Data with an entry for each data row of the table. Each entry should have data for
1001                           each column in the row.
1002        :param include_header: If True, then a header will be included at top of table. (Defaults to True)
1003        """
1004        table_buf = io.StringIO()
1005
1006        if include_header:
1007            header = self.generate_header()
1008            table_buf.write(header)
1009        else:
1010            top_border = self.generate_table_top_border()
1011            table_buf.write(top_border)
1012
1013        table_buf.write('\n')
1014
1015        for index, row_data in enumerate(table_data):
1016            if index > 0:
1017                row_bottom_border = self.generate_row_bottom_border()
1018                table_buf.write(row_bottom_border)
1019                table_buf.write('\n')
1020
1021            row = self.generate_data_row(row_data)
1022            table_buf.write(row)
1023            table_buf.write('\n')
1024
1025        table_buf.write(self.generate_table_bottom_border())
1026        return table_buf.getvalue()
1027
1028
1029class AlternatingTable(BorderedTable):
1030    """
1031    Implementation of BorderedTable which uses background colors to distinguish between rows instead of row border
1032    lines. This class can be used to create the whole table at once or one row at a time.
1033
1034    To nest an AlternatingTable within another AlternatingTable, set style_data_text to False on the Column
1035    which contains the nested table. That will prevent the current row's background color from affecting the colors
1036    of the nested table.
1037    """
1038
1039    def __init__(
1040        self,
1041        cols: Sequence[Column],
1042        *,
1043        tab_width: int = 4,
1044        column_borders: bool = True,
1045        padding: int = 1,
1046        border_fg: Optional[ansi.FgColor] = None,
1047        border_bg: Optional[ansi.BgColor] = None,
1048        header_bg: Optional[ansi.BgColor] = None,
1049        odd_bg: Optional[ansi.BgColor] = None,
1050        even_bg: Optional[ansi.BgColor] = ansi.Bg.DARK_GRAY,
1051    ) -> None:
1052        """
1053        AlternatingTable initializer
1054
1055        Note: Specify background colors using subclasses of BgColor (e.g. Bg, EightBitBg, RgbBg)
1056
1057        :param cols: column definitions for this table
1058        :param tab_width: all tabs will be replaced with this many spaces. If a row's fill_char is a tab,
1059                          then it will be converted to one space.
1060        :param column_borders: if True, borders between columns will be included. This gives the table a grid-like
1061                               appearance. Turning off column borders results in a unified appearance between
1062                               a row's cells. (Defaults to True)
1063        :param padding: number of spaces between text and left/right borders of cell
1064        :param border_fg: optional foreground color for borders (defaults to None)
1065        :param border_bg: optional background color for borders (defaults to None)
1066        :param header_bg: optional background color for header cells (defaults to None)
1067        :param odd_bg: optional background color for odd numbered data rows (defaults to None)
1068        :param even_bg: optional background color for even numbered data rows (defaults to StdBg.DARK_GRAY)
1069        :raises: ValueError if tab_width is less than 1
1070        :raises: ValueError if padding is less than 0
1071        """
1072        super().__init__(
1073            cols,
1074            tab_width=tab_width,
1075            column_borders=column_borders,
1076            padding=padding,
1077            border_fg=border_fg,
1078            border_bg=border_bg,
1079            header_bg=header_bg,
1080        )
1081        self.row_num = 1
1082        self.odd_bg = odd_bg
1083        self.even_bg = even_bg
1084
1085    def apply_data_bg(self, value: Any) -> str:
1086        """
1087        Apply background color to data text based on what row is being generated and whether a color has been defined
1088        :param value: object whose text is to be colored
1089        :return: formatted data string
1090        """
1091        if self.row_num % 2 == 0 and self.even_bg is not None:
1092            return ansi.style(value, bg=self.even_bg)
1093        elif self.row_num % 2 != 0 and self.odd_bg is not None:
1094            return ansi.style(value, bg=self.odd_bg)
1095        else:
1096            return str(value)
1097
1098    def generate_data_row(self, row_data: Sequence[Any]) -> str:
1099        """
1100        Generate a data row
1101
1102        :param row_data: data with an entry for each column in the row
1103        :return: data row string
1104        """
1105        row = super().generate_data_row(row_data)
1106        self.row_num += 1
1107        return row
1108
1109    def generate_table(self, table_data: Sequence[Sequence[Any]], *, include_header: bool = True) -> str:
1110        """
1111        Generate a table from a data set
1112
1113        :param table_data: Data with an entry for each data row of the table. Each entry should have data for
1114                           each column in the row.
1115        :param include_header: If True, then a header will be included at top of table. (Defaults to True)
1116        """
1117        table_buf = io.StringIO()
1118
1119        if include_header:
1120            header = self.generate_header()
1121            table_buf.write(header)
1122        else:
1123            top_border = self.generate_table_top_border()
1124            table_buf.write(top_border)
1125
1126        table_buf.write('\n')
1127
1128        for row_data in table_data:
1129            row = self.generate_data_row(row_data)
1130            table_buf.write(row)
1131            table_buf.write('\n')
1132
1133        table_buf.write(self.generate_table_bottom_border())
1134        return table_buf.getvalue()
1135