sqlmesh.utils.rich

  1from __future__ import annotations
  2
  3import typing as t
  4
  5import re
  6
  7from rich.console import Console
  8from rich.progress import Column, ProgressColumn, Task, Text
  9from rich.theme import Theme
 10from rich.table import Table
 11from rich.align import Align
 12
 13if t.TYPE_CHECKING:
 14    import pandas as pd
 15
 16theme = Theme(
 17    {
 18        "added": "green",
 19        "removed": "red",
 20        "direct": "magenta",  # directly modified
 21        "indirect": "yellow",  # indirectly modified
 22        "metadata": "cyan",  # metadata updated
 23    }
 24)
 25
 26console = Console(theme=theme)
 27
 28
 29class BatchColumn(ProgressColumn):
 30    """Renders completed count/total, "pending".
 31
 32    Space pads the completed count so that progress length does not change as task progresses
 33    past powers of 10.
 34
 35    Source: https://rich.readthedocs.io/en/stable/reference/progress.html#rich.progress.MofNCompleteColumn
 36
 37    Args:
 38        separator: Text to separate completed and total values. Defaults to "/".
 39    """
 40
 41    def __init__(self, separator: str = "/", table_column: t.Optional[Column] = None):
 42        self.separator = separator
 43        super().__init__(table_column=table_column)
 44
 45    def render(self, task: Task) -> Text:
 46        """Show completed count/total, "pending"."""
 47        total = int(task.total) if task.total is not None else "?"
 48        completed = int(task.completed)
 49
 50        if completed == 0 and task.total is not None and task.total > 0:
 51            return Text("pending", style="progress.download")
 52
 53        total_width = len(str(total))
 54        return Text(
 55            f"{completed:{total_width}d}{self.separator}{total}",
 56            style="progress.download",
 57        )
 58
 59
 60def strip_ansi_codes(text: str) -> str:
 61    """Strip ANSI color codes and styling from text."""
 62    ansi_escape = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]")
 63    return ansi_escape.sub("", text).strip()
 64
 65
 66def df_to_table(
 67    header: str,
 68    df: pd.DataFrame,
 69    show_index: bool = True,
 70    index_name: str = "Row",
 71) -> Table:
 72    """Convert a pandas.DataFrame obj into a rich.Table obj.
 73    Args:
 74        df (DataFrame): A Pandas DataFrame to be converted to a rich Table.
 75        rich_table (Table): A rich Table that should be populated by the DataFrame values.
 76        show_index (bool): Add a column with a row count to the table. Defaults to True.
 77        index_name (str, optional): The column name to give to the index column. Defaults to None, showing no value.
 78    Returns:
 79        Table: The rich Table instance passed, populated with the DataFrame values."""
 80
 81    rich_table = Table(title=f"[bold red]{header}[/bold red]", show_lines=True, min_width=60)
 82    if show_index:
 83        index_name = str(index_name) if index_name else ""
 84        rich_table.add_column(Align.center(index_name))
 85
 86    for column in df.columns:
 87        column_name = column if isinstance(column, str) else ": ".join(str(col) for col in column)
 88
 89        # Color coding unit test columns (expected/actual), can be removed or refactored if df_to_table is used elswhere too
 90        lower = column_name.lower()
 91        if "expected" in lower:
 92            column_name = f"[green]{column_name}[/green]"
 93        elif "actual" in lower:
 94            column_name = f"[red]{column_name}[/red]"
 95
 96        rich_table.add_column(Align.center(column_name))
 97
 98    for index, value_list in zip(df.index, df.values.tolist()):
 99        row = [str(index)] if show_index else []
100        row += [str(x) for x in value_list]
101        center = [Align.center(x) for x in row]
102        rich_table.add_row(*center)
103
104    return rich_table