summaryrefslogtreecommitdiff
path: root/lib/spack/llnl/util/tty/colify.py
blob: ca157e3303756f6e6f44cc4da9c1e551fd40c2f4 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
# Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

"""
Routines for printing columnar output.  See ``colify()`` for more information.
"""
from __future__ import division, unicode_literals

import os
import sys

from six import StringIO, text_type

from llnl.util.tty import terminal_size
from llnl.util.tty.color import cextra, clen


class ColumnConfig:

    def __init__(self, cols):
        self.cols = cols
        self.line_length = 0
        self.valid = True
        self.widths = [0] * cols   # does not include ansi colors

    def __repr__(self):
        attrs = [(a, getattr(self, a))
                 for a in dir(self) if not a.startswith("__")]
        return "<Config: %s>" % ", ".join("%s: %r" % a for a in attrs)


def config_variable_cols(elts, console_width, padding, cols=0):
    """Variable-width column fitting algorithm.

       This function determines the most columns that can fit in the
       screen width.  Unlike uniform fitting, where all columns take
       the width of the longest element in the list, each column takes
       the width of its own longest element. This packs elements more
       efficiently on screen.

       If cols is nonzero, force
    """
    if cols < 0:
        raise ValueError("cols must be non-negative.")

    # Get a bound on the most columns we could possibly have.
    # 'clen' ignores length of ansi color sequences.
    lengths = [clen(e) for e in elts]
    max_cols = max(1, console_width // (min(lengths) + padding))
    max_cols = min(len(elts), max_cols)

    # Range of column counts to try.  If forced, use the supplied value.
    col_range = [cols] if cols else range(1, max_cols + 1)

    # Determine the most columns possible for the console width.
    configs = [ColumnConfig(c) for c in col_range]
    for i, length in enumerate(lengths):
        for conf in configs:
            if conf.valid:
                col = i // ((len(elts) + conf.cols - 1) // conf.cols)
                p = padding if col < (conf.cols - 1) else 0

                if conf.widths[col] < (length + p):
                    conf.line_length += length + p - conf.widths[col]
                    conf.widths[col]  = length + p
                    conf.valid = (conf.line_length < console_width)

    try:
        config = next(conf for conf in reversed(configs) if conf.valid)
    except StopIteration:
        # If nothing was valid the screen was too narrow -- just use 1 col.
        config = configs[0]

    config.widths = [w for w in config.widths if w != 0]
    config.cols = len(config.widths)
    return config


def config_uniform_cols(elts, console_width, padding, cols=0):
    """Uniform-width column fitting algorithm.

       Determines the longest element in the list, and determines how
       many columns of that width will fit on screen.  Returns a
       corresponding column config.
    """
    if cols < 0:
        raise ValueError("cols must be non-negative.")

    # 'clen' ignores length of ansi color sequences.
    max_len = max(clen(e) for e in elts) + padding
    if cols == 0:
        cols = max(1, console_width // max_len)
        cols = min(len(elts), cols)

    config = ColumnConfig(cols)
    config.widths = [max_len] * cols

    return config


def colify(elts, **options):
    """Takes a list of elements as input and finds a good columnization
    of them, similar to how gnu ls does. This supports both
    uniform-width and variable-width (tighter) columns.

    If elts is not a list of strings, each element is first conveted
    using ``str()``.

    Keyword Arguments:
        output (stream): A file object to write to. Default is ``sys.stdout``
        indent (int):    Optionally indent all columns by some number of spaces
        padding (int):   Spaces between columns. Default is 2
        width (int):     Width of the output. Default is 80 if tty not detected
        cols (int):      Force number of columns. Default is to size to
                         terminal, or single-column if no tty
        tty (bool):      Whether to attempt to write to a tty. Default is to
                         autodetect a tty. Set to False to force single-column
                         output
        method (str):    Method to use to fit columns. Options are variable or
                         uniform. Variable-width columns are tighter, uniform
                         columns are all the same width and fit less data on
                         the screen
    """
    # Get keyword arguments or set defaults
    cols         = options.pop("cols", 0)
    output       = options.pop("output", sys.stdout)
    indent       = options.pop("indent", 0)
    padding      = options.pop("padding", 2)
    tty          = options.pop('tty', None)
    method       = options.pop("method", "variable")
    console_cols = options.pop("width", None)

    if options:
        raise TypeError(
            "'%s' is an invalid keyword argument for this function."
            % next(options.iterkeys()))

    # elts needs to be an array of strings so we can count the elements
    elts = [text_type(elt) for elt in elts]
    if not elts:
        return (0, ())

    # environment size is of the form "<rows>x<cols>"
    env_size = os.environ.get('COLIFY_SIZE')
    if env_size:
        try:
            r, c = env_size.split('x')
            console_rows, console_cols = int(r), int(c)
            tty = True
        except BaseException:
            pass

    # Use only one column if not a tty.
    if not tty:
        if tty is False or not output.isatty():
            cols = 1

    # Specify the number of character columns to use.
    if not console_cols:
        console_rows, console_cols = terminal_size()
    elif type(console_cols) != int:
        raise ValueError("Number of columns must be an int")
    console_cols = max(1, console_cols - indent)

    # Choose a method.  Variable-width colums vs uniform-width.
    if method == "variable":
        config = config_variable_cols(elts, console_cols, padding, cols)
    elif method == "uniform":
        config = config_uniform_cols(elts, console_cols, padding, cols)
    else:
        raise ValueError("method must be either 'variable' or 'uniform'")

    cols = config.cols
    rows = (len(elts) + cols - 1) // cols
    rows_last_col = len(elts) % rows

    for row in range(rows):
        output.write(" " * indent)
        for col in range(cols):
            elt = col * rows + row
            width = config.widths[col] + cextra(elts[elt])
            if col < cols - 1:
                fmt = '%%-%ds' % width
                output.write(fmt % elts[elt])
            else:
                # Don't pad the rightmost column (sapces can wrap on
                # small teriminals if one line is overlong)
                output.write(elts[elt])

        output.write("\n")
        row += 1
        if row == rows_last_col:
            cols -= 1

    return (config.cols, tuple(config.widths))


def colify_table(table, **options):
    """Version of ``colify()`` for data expressed in rows, (list of lists).

       Same as regular colify but:

       1. This takes a list of lists, where each sub-list must be the
          same length, and each is interpreted as a row in a table.
          Regular colify displays a sequential list of values in columns.

       2. Regular colify will always print with 1 column when the output
          is not a tty.  This will always print with same dimensions of
          the table argument.

    """
    if table is None:
        raise TypeError("Can't call colify_table on NoneType")
    elif not table or not table[0]:
        raise ValueError("Table is empty in colify_table!")

    columns = len(table[0])

    def transpose():
        for i in range(columns):
            for row in table:
                yield row[i]

    if 'cols' in options:
        raise ValueError("Cannot override columsn in colify_table.")
    options['cols'] = columns

    # don't reduce to 1 column for non-tty
    options['tty'] = True

    colify(transpose(), **options)


def colified(elts, **options):
    """Invokes the ``colify()`` function but returns the result as a string
       instead of writing it to an output string."""
    sio = StringIO()
    options['output'] = sio
    colify(elts, **options)
    return sio.getvalue()