diff --git a/dvc/utils/table.py b/dvc/utils/table.py index 6d997549a..55bf93b1c 100644 --- a/dvc/utils/table.py +++ b/dvc/utils/table.py @@ -1,60 +1,17 @@ -from dataclasses import dataclass -from typing import TYPE_CHECKING, List, cast +from typing import TYPE_CHECKING, Any, List -from rich.style import StyleType -from rich.table import Column as RichColumn from rich.table import Table as RichTable if TYPE_CHECKING: - from rich.console import ( - Console, - ConsoleOptions, - JustifyMethod, - OverflowMethod, - RenderableType, - ) - - -@dataclass -class Column(RichColumn): - collapse: bool = False + from rich.console import Console, ConsoleOptions class Table(RichTable): - def add_column( # type: ignore[override] # pylint: disable=arguments-differ,arguments-renamed # noqa: E501 - self, - header: "RenderableType" = "", - footer: "RenderableType" = "", - *, - header_style: StyleType = None, - footer_style: StyleType = None, - style: StyleType = None, - justify: "JustifyMethod" = "left", - overflow: "OverflowMethod" = "ellipsis", - width: int = None, - min_width: int = None, - max_width: int = None, - ratio: int = None, - no_wrap: bool = False, - collapse: bool = False, + def add_column( # pylint: disable=arguments-differ + self, *args: Any, collapse: bool = False, **kwargs: Any ) -> None: - column = Column( # type: ignore[call-arg] - _index=len(self.columns), - header=header, - footer=footer, - header_style=header_style or "", - footer_style=footer_style or "", - style=style or "", - justify=justify, - overflow=overflow, - width=width, - min_width=min_width, - max_width=max_width, - ratio=ratio, - no_wrap=no_wrap, - collapse=collapse, - ) - self.columns.append(column) + super().add_column(*args, **kwargs) + self.columns[-1].collapse = collapse # type: ignore[attr-defined] def _calculate_column_widths( self, console: "Console", options: "ConsoleOptions" @@ -67,9 +24,12 @@ class Table(RichTable): """ widths = super()._calculate_column_widths(console, options) last_collapsed = -1 - columns = cast(List[Column], self.columns) + columns = self.columns for i in range(len(columns) - 1, -1, -1): - if widths[i] == 0 and columns[i].collapse: + if ( + widths[i] == 0 + and columns[i].collapse # type: ignore[attr-defined] + ): if last_collapsed >= 0: del widths[last_collapsed] del columns[last_collapsed] @@ -99,14 +59,15 @@ class Table(RichTable): If table is still too wide after collapsing, rich's automatic overflow handling will be used. """ - columns = cast(List[Column], self.columns) - collapsible = [column.collapse for column in columns] + 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]: - total_width -= widths[i] excess_width -= widths[i] widths[i] = 0 if excess_width <= 0: