iterative/dvc

from typing import TYPE_CHECKING, Any, List

from rich.table import Table as RichTable

if TYPE_CHECKING:
    from rich.console import Console, ConsoleOptions


class Table(RichTable):
    def add_column(  # pylint: disable=arguments-differ
        self, *args: Any, collapse: bool = False, **kwargs: Any
    ) -> None:
        super().add_column(*args, **kwargs)
        self.columns[-1].collapse = collapse  # type: ignore[attr-defined]

    def _calculate_column_widths(
        self, console: "Console", options: "ConsoleOptions"
    ) -> List[int]:
        """Calculate the widths of each column, including padding, not
        including borders.

        Adjacent collapsed columns will be removed until there is only a single
        truncated column remaining.
        """
        widths = super()._calculate_column_widths(console, options)
        last_collapsed = -1
        columns = self.columns
        for i in range(len(columns) - 1, -1, -1):
            if (
                widths[i] == 0
                and columns[i].collapse  # type: ignore[attr-defined]
            ):
                if last_collapsed >= 0:
                    del widths[last_collapsed]
                    del columns[last_collapsed]
                    if self.box:
                        options.max_width += 1
                    for column in columns[last_collapsed:]:
                        column._index -= 1
                last_collapsed = i
                padding = self._get_padding_width(i)
                if (
                    columns[i].overflow == "ellipsis"
                    and (sum(widths) + padding) <= options.max_width
                ):
                    # Set content width to 1 (plus padding) if we can fit a
                    # single unicode ellipsis in this column
                    widths[i] = 1 + padding
            else:
                last_collapsed = -1
        return widths

    def _collapse_widths(  # type: ignore[override]
        # pylint: disable=arguments-differ
        self,
        widths: List[int],
        wrapable: List[bool],
        max_width: int,
    ) -> List[int]:
        """Collapse columns right-to-left if possible to fit table into
        max_width.

        If table is still too wide after collapsing, rich's automatic overflow
        handling will be used.
        """
        collapsible = [
            column.collapse  # type: ignore[attr-defined]
            for column in self.columns
        ]
        total_width = sum(widths)
        excess_width = total_width - max_width
        if any(collapsible):
            for i in range(len(widths) - 1, -1, -1):
                if collapsible[i]:
                    excess_width -= widths[i]
                    widths[i] = 0
                    if excess_width <= 0:
                        break
        return super()._collapse_widths(widths, wrapable, max_width)