1from __future__ import division
2
3import textwrap
4
5from typing import List
6
7from clikit.api.formatter import Formatter
8from clikit.utils.string import get_max_line_length
9from clikit.utils.string import get_max_word_length
10from clikit.utils.string import get_string_length
11
12
13class CellWrapper:
14    """
15    Wraps cells to fit a given screen width with a given number of columns.
16    """
17
18    def __init__(self):  # type: () -> None
19        self._cells = []
20        self._cell_lengths = []
21        self._wrapped_rows = []
22        self._nb_columns = 0
23        self._column_lengths = []
24        self._word_wraps = False
25        self._word_cuts = False
26        self._max_total_width = 0
27        self._total_width = 0
28
29    @property
30    def cells(self):  # type: () -> List[str]
31        return self._cells
32
33    @property
34    def wrapped_rows(self):  # type: () -> List[List[str]]
35        return self._wrapped_rows
36
37    @property
38    def column_lengths(self):  # type: () -> List[int]
39        return self._column_lengths
40
41    @property
42    def nb_columns(self):  # type: () -> int
43        return self._nb_columns
44
45    @property
46    def max_total_width(self):  # type: () -> int
47        return self._max_total_width
48
49    @property
50    def total_width(self):  # type: () -> int
51        return self._total_width
52
53    def add_cell(self, cell):  # type: (str) -> CellWrapper
54        self._cells.append(cell.rstrip())
55
56        return self
57
58    def add_cells(self, cells):  # type: (List[str]) -> CellWrapper
59        for cell in cells:
60            self.add_cell(cell)
61
62        return self
63
64    def get_estimated_nb_columns(self, max_total_width):  # type: (int) -> int
65        """
66        Returns an estimated number of columns for the given maximum width.
67        """
68        row_width = 0
69
70        for i, cell in enumerate(self._cells):
71            row_width += get_string_length(cell)
72
73            if row_width > max_total_width:
74                return i
75
76        return len(self._cells) - 1
77
78    def has_word_wraps(self):  # type: () -> bool
79        return self._word_wraps
80
81    def has_word_cuts(self):  # type: () -> bool
82        return self._word_cuts
83
84    def fit(
85        self, max_total_width, nb_columns, formatter
86    ):  # type: (int, int, Formatter) -> None
87        self._reset_state(max_total_width, nb_columns)
88        self._init_rows(formatter)
89
90        # If the cells fit within the max width we're good
91        if self._total_width <= max_total_width:
92            return
93
94        self._wrap_columns(formatter)
95
96    def _reset_state(self, max_total_width, nb_columns):  # type: (int, int) -> None
97        self._wrapped_rows = []
98        self._nb_columns = nb_columns
99        self._cell_lengths = []
100        self._column_lengths = [0] * nb_columns
101        self._word_wraps = False
102        self._word_cuts = False
103        self._max_total_width = max_total_width
104        self._total_width = 0
105
106    def _init_rows(self, formatter):  # type: (Formatter) -> None
107        col = 0
108
109        for i, cell in enumerate(self._cells):
110            if col == 0:
111                self._wrapped_rows.append([""] * self._nb_columns)
112                self._cell_lengths.append([0] * self._nb_columns)
113
114            self._wrapped_rows[-1][col] = cell
115            self._cell_lengths[-1][col] = get_string_length(cell, formatter)
116            self._column_lengths[col] = max(
117                self._column_lengths[col], self._cell_lengths[-1][col]
118            )
119
120            col = (col + 1) % self._nb_columns
121
122        # Fill last row
123        if col > 0:
124            while col < self._nb_columns:
125                self._wrapped_rows[-1][col] = ""
126                self._cell_lengths[-1][col] = 0
127                col += 1
128
129        self._total_width = sum(self._column_lengths)
130
131    def _wrap_columns(self, formatter):  # type: (Formatter) -> None
132        available_width = self._max_total_width
133        long_column_lengths = self._column_lengths[:]
134
135        # Filter "short" column, i.e. columns that are not wrapped
136        # We distribute the available screen width by the number of columns
137        # and decide that all columns that are shorter than their share are
138        # "short".
139        # This process is repeated until no more "short" columns are found.
140        repeat = True
141        while repeat:
142            threshold = available_width / len(long_column_lengths)
143            repeat = False
144
145            for col, length in enumerate(long_column_lengths):
146                if length is not None and length <= threshold:
147                    available_width -= length
148                    long_column_lengths[col] = None
149                    repeat = True
150
151        # Calculate actual and available width
152        actual_width = 0
153        last_adapted_col = 0
154
155        # "Long" columns, i.e. columns that need to be wrapped, are added to
156        # the actual width
157        for col, length in enumerate(long_column_lengths):
158            if length is None:
159                continue
160
161            actual_width += length
162            last_adapted_col = col
163
164        # Fit columns into available width
165        for col, length in enumerate(long_column_lengths):
166            if length is None:
167                continue
168
169            # Keep ratios of column lengths and distribute them among the
170            # available width
171            self._column_lengths[col] = int(
172                round((length / actual_width) * available_width)
173            )
174
175            if col == last_adapted_col:
176                # Fix rounding errors
177                self._column_lengths[col] += self._max_total_width - sum(
178                    self._column_lengths
179                )
180
181            self._wrap_column(col, self._column_lengths[col], formatter)
182
183            # Recalculate the column length based on the actual wrapped length
184            self._refresh_column_length(col)
185
186            # Recalculate the actual width based on the changed length.
187            actual_width = actual_width - length + self._column_lengths[col]
188
189        self._total_width = sum(self._column_lengths)
190
191    def _wrap_column(
192        self, col, column_length, formatter
193    ):  # type: (int, List[int], Formatter) -> None
194        for i, row in enumerate(self._wrapped_rows):
195            cell = row[col]
196            cell_length = self._cell_lengths[i][col]
197
198            if cell_length > column_length:
199                self._word_wraps = True
200
201                if not self._word_cuts:
202                    min_length_without_cut = get_max_word_length(cell, formatter)
203
204                    if min_length_without_cut > column_length:
205                        self._word_cuts = True
206
207                # TODO: use format aware wrapper
208                wrapped_cell = "\n".join(textwrap.wrap(cell, column_length))
209
210                self._wrapped_rows[i][col] = wrapped_cell
211
212                # Refresh cell length
213                self._cell_lengths[i][col] = get_max_line_length(
214                    wrapped_cell, formatter
215                )
216
217    def _refresh_column_length(self, col):  # type: (int) -> None
218        self._column_lengths[col] = 0
219
220        for i, row in enumerate(self._wrapped_rows):
221            self._column_lengths[col] = max(
222                self._column_lengths[col], self._cell_lengths[i][col]
223            )
224