|
| 1 | +# Copyright (C) British Crown (Met Office) & Contributors. |
| 2 | +# This file is part of Rose, a framework for meteorological suites. |
| 3 | +# |
| 4 | +# Rose is free software: you can redistribute it and/or modify |
| 5 | +# it under the terms of the GNU General Public License as published by |
| 6 | +# the Free Software Foundation, either version 3 of the License, or |
| 7 | +# (at your option) any later version. |
| 8 | +# |
| 9 | +# Rose is distributed in the hope that it will be useful, |
| 10 | +# but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 11 | +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 12 | +# GNU General Public License for more details. |
| 13 | +# |
| 14 | +# You should have received a copy of the GNU General Public License |
| 15 | +# along with Rose. If not, see <http://www.gnu.org/licenses/>. |
| 16 | +# ----------------------------------------------------------------------------- |
| 17 | + |
| 18 | +from subprocess import Popen, PIPE |
| 19 | +import textwrap |
| 20 | +from typing import Optional, List |
| 21 | + |
| 22 | +# NOTE: Box Drawing characters |
| 23 | +# t, b, l, r -> top, bottom, left, right |
| 24 | +# h, v -> horizontal, vertical |
| 25 | +# l, x -> line, cross (i.e. vertex) |
| 26 | + |
| 27 | +# unicode (U_) chars |
| 28 | +U_HL = '─' |
| 29 | +U_VL = '│' |
| 30 | +U_TL = '┌' |
| 31 | +U_TR = '┐' |
| 32 | +U_BL = '└' |
| 33 | +U_BR = '┘' |
| 34 | +U_LX = '├' |
| 35 | +U_RX = '┤' |
| 36 | +U_TX = '┬' |
| 37 | +U_BX = '┴' |
| 38 | +U_XX = '┼' |
| 39 | + |
| 40 | +# ASCII (A_) chars |
| 41 | +A_HL = '-' |
| 42 | +A_VL = '|' |
| 43 | +A_TL = '+' |
| 44 | +A_TR = '+' |
| 45 | +A_BL = '+' |
| 46 | +A_BR = '+' |
| 47 | +A_LX = '+' |
| 48 | +A_RX = '+' |
| 49 | +A_TX = '+' |
| 50 | +A_BX = '+' |
| 51 | +A_XX = '+' |
| 52 | + |
| 53 | + |
| 54 | +def get_terminal_cols(default=80): |
| 55 | + """Return the width of the terminal in columns. |
| 56 | +
|
| 57 | + If the check fails, this will fallback to the default value. |
| 58 | + """ |
| 59 | + proc = Popen(['stty', 'size'], stdout=PIPE, stderr=PIPE, text=True) |
| 60 | + try: |
| 61 | + return int(proc.communicate()[0].split()[1]) or default |
| 62 | + except Exception: |
| 63 | + return default |
| 64 | + |
| 65 | + |
| 66 | +def table( |
| 67 | + rows: List[List[str]], |
| 68 | + header: Optional[List[str]] = None, |
| 69 | + max_width: Optional[int] = None, |
| 70 | + unicode: bool = True, |
| 71 | +) -> str: |
| 72 | + """Format text into a table. |
| 73 | +
|
| 74 | + Args: |
| 75 | + rows: 2D table composed of lists. |
| 76 | + header: Optional 1D list to use as the table header. |
| 77 | + max_width: Maximum permitted width for the whole table including |
| 78 | + borders. |
| 79 | + unicode: Use unicode characters if True, else fallback to ascii. |
| 80 | +
|
| 81 | + Examples: |
| 82 | + # simple ASCII table: |
| 83 | + >>> print(table([['a', 'b', 'c']], unicode=False)) |
| 84 | + +---+---+---+ |
| 85 | + | a | b | c | |
| 86 | + +---+---+---+ |
| 87 | +
|
| 88 | + # unicode table with headers and a max-width: |
| 89 | + >>> print(table( |
| 90 | + ... [['a', 'b', 'c'], ['d', 'e', 'qwertyuiopasdfghjkl']], |
| 91 | + ... header=['foo', 'bar', 'baz'], |
| 92 | + ... max_width=27, |
| 93 | + ... )) |
| 94 | + ┌─────┬─────┬─────────────┐ |
| 95 | + │ foo │ bar │ baz │ |
| 96 | + ├─────┼─────┼─────────────┤ |
| 97 | + ├─────┼─────┼─────────────┤ |
| 98 | + │ a │ b │ c │ |
| 99 | + ├─────┼─────┼─────────────┤ |
| 100 | + │ d │ e │ qwertyuiopa │ |
| 101 | + │ │ │ sdfghjkl │ |
| 102 | + └─────┴─────┴─────────────┘ |
| 103 | +
|
| 104 | + # edge case: table doesn't fit within max_width |
| 105 | + # (columns will use width of 1 and table will exceed max_width) |
| 106 | + >>> print(table([["a", "long"]], max_width=6)) |
| 107 | + ┌───┬───┐ |
| 108 | + │ a │ l │ |
| 109 | + │ │ o │ |
| 110 | + │ │ n │ |
| 111 | + │ │ g │ |
| 112 | + └───┴───┘ |
| 113 | +
|
| 114 | + """ |
| 115 | + # determine character set |
| 116 | + if unicode: |
| 117 | + hl = U_HL |
| 118 | + vl = U_VL |
| 119 | + tl = U_TL |
| 120 | + tr = U_TR |
| 121 | + bl = U_BL |
| 122 | + br = U_BR |
| 123 | + lx = U_LX |
| 124 | + rx = U_RX |
| 125 | + tx = U_TX |
| 126 | + bx = U_BX |
| 127 | + xx = U_XX |
| 128 | + else: |
| 129 | + hl = A_HL |
| 130 | + vl = A_VL |
| 131 | + tl = A_TL |
| 132 | + tr = A_TR |
| 133 | + bl = A_BL |
| 134 | + br = A_BR |
| 135 | + lx = A_LX |
| 136 | + rx = A_RX |
| 137 | + tx = A_TX |
| 138 | + bx = A_BX |
| 139 | + xx = A_XX |
| 140 | + |
| 141 | + # determine column widths |
| 142 | + widths = [0] * len(rows[0]) |
| 143 | + for table_ in (rows, [header or []]): |
| 144 | + for row in (table_): |
| 145 | + for ind, col in enumerate(row): |
| 146 | + widths[ind] = max(widths[ind], len(col)) |
| 147 | + |
| 148 | + def calc_width(): |
| 149 | + return (sum(widths) + (len(widths) * 3) + 1) |
| 150 | + |
| 151 | + # resize cols to make them fix the max_width |
| 152 | + if max_width: |
| 153 | + overhang = max_width - calc_width() |
| 154 | + for _ in range(15): # limit to 15 itterations |
| 155 | + _max_width = 0 |
| 156 | + _max_col = 0 |
| 157 | + # find the longest column |
| 158 | + for ind, width in enumerate(widths): |
| 159 | + if width > _max_width: |
| 160 | + _max_width = width |
| 161 | + _max_col = ind |
| 162 | + # reduce the longest column by up to 50% until it fits |
| 163 | + widths[_max_col] = max( |
| 164 | + int(widths[_max_col] / 2), |
| 165 | + widths[_max_col] + overhang |
| 166 | + ) or 1 # don't let column width go to 0 |
| 167 | + overhang = max_width - calc_width() |
| 168 | + if overhang >= 0: |
| 169 | + break |
| 170 | + |
| 171 | + # textwrap the table to fit the desired widths |
| 172 | + _rows = [ |
| 173 | + [ |
| 174 | + textwrap.wrap(col, width=widths[ind]) |
| 175 | + for ind, col in enumerate(row) |
| 176 | + ] |
| 177 | + for row in rows |
| 178 | + ] |
| 179 | + |
| 180 | + # Python f-string for formatting each row of the table |
| 181 | + row_format = vl + vl.join( |
| 182 | + f" {{{ind}:{width}}} " for ind, width in enumerate(widths) |
| 183 | + ) + vl |
| 184 | + |
| 185 | + # a divider row (i.e. a horizontal line) |
| 186 | + blank_line = lx + xx.join(hl * (width + 2) for width in widths) + rx |
| 187 | + |
| 188 | + # top border of table |
| 189 | + ret = [tl + tx.join(hl * (width + 2) for width in widths) + tr] |
| 190 | + |
| 191 | + # table header |
| 192 | + if header: |
| 193 | + ret.append(row_format.format(*header)) |
| 194 | + ret.extend([blank_line, blank_line]) |
| 195 | + |
| 196 | + # table body |
| 197 | + for row_ind, row in enumerate(_rows): |
| 198 | + max_height = max(len(col) for col in row) |
| 199 | + for line_ind in range(max_height): |
| 200 | + ret.append(row_format.format( |
| 201 | + *[ |
| 202 | + col[line_ind] if line_ind < len(col) else '' |
| 203 | + for col in row |
| 204 | + ] |
| 205 | + )) |
| 206 | + if row_ind != len(_rows) - 1: |
| 207 | + # don't add a divider row for the last item |
| 208 | + ret.append(blank_line) |
| 209 | + |
| 210 | + # bottom border of table |
| 211 | + ret.append(bl + bx.join(hl * (width + 2) for width in widths) + br) |
| 212 | + |
| 213 | + return '\n'.join(ret) |
0 commit comments