feat(ux): use rich to format Table.info() output

Author: cpcloud

ibis/backends/tests/test_generic.py

@@ -2,6 +2,7 @@
 import importlib
 import io
 import operator
+from contextlib import redirect_stdout
 
 import numpy as np
 import pandas as pd
@@ -435,16 +436,25 @@ def test_select_sort_sort(alltypes):
     query = query.sort_by(query.year).sort_by(query.bool_col)
 
 
-def test_table_info(alltypes):
+def check_table_info(buf, schema):
+    info_str = buf.getvalue()
+
+    assert "Null" in info_str
+    assert all(type.__class__.__name__ in info_str for type in schema.types)
+    assert all(name in info_str for name in schema.names)
+
+
+def test_table_info_buf(alltypes):
     buf = io.StringIO()
     alltypes.info(buf=buf)
+    check_table_info(buf, alltypes.schema())
 
-    info_str = buf.getvalue()
-    schema = alltypes.schema()
 
-    assert "Nulls" in info_str
-    assert all(str(type) in info_str for type in schema.types)
-    assert all(name in info_str for name in schema.names)
+def test_table_info_no_buf(alltypes):
+    buf = io.StringIO()
+    with redirect_stdout(buf):
+        alltypes.info()
+    check_table_info(buf, alltypes.schema())
 
 
 @pytest.mark.parametrize(

ibis/common/grounds.py

@@ -5,12 +5,16 @@
 from typing import Any, Hashable
 from weakref import WeakValueDictionary
 
+from rich.console import Console
+
 from ibis.common.caching import WeakCache
 from ibis.common.validators import ImmutableProperty, Optional, Validator
 from ibis.util import frozendict
 
 EMPTY = inspect.Parameter.empty  # marker for missing argument
 
+console = Console()
+
 
 class BaseMeta(ABCMeta):
 

ibis/expr/types/relations.py

@@ -4,17 +4,20 @@
 import functools
 import itertools
 import operator
+import sys
 import warnings
 from functools import cached_property
 from typing import IO, TYPE_CHECKING, Any, Iterable, Literal, Mapping, Sequence
 
 import numpy as np
-import tabulate
+import rich.pretty
+import rich.table
 from public import public
 
 import ibis
 from ibis import util
 from ibis.common import exceptions as com
+from ibis.common.grounds import console
 from ibis.expr.deferred import Deferred
 from ibis.expr.types.core import Expr
 
@@ -872,29 +875,37 @@ def info(self, buf: IO[str] | None = None) -> None:
         buf
             A writable buffer, defaults to stdout
         """
+        if buf is None:
+            buf = sys.stdout
+
         metrics = [self[col].count().name(col) for col in self.columns]
         metrics.append(self.count().name("nrows"))
 
         schema = self.schema()
 
         *items, (_, n) = self.aggregate(metrics).execute().squeeze().items()
 
-        tabulated = tabulate.tabulate(
-            [
-                (
-                    column,
-                    schema[column],
-                    f"{n - non_nulls} ({100 * (1.0 - non_nulls / n):>3.3g}%)",
-                )
-                for column, non_nulls in items
-            ],
-            headers=["Column", "Type", "Nulls (%)"],
-            colalign=("left", "left", "right"),
-        )
-        width = tabulated[tabulated.index("\n") + 1 :].index("\n")
-        row_count = f"Rows: {n}".center(width)
-        footer_line = "-" * width
-        print("\n".join([tabulated, footer_line, row_count]), file=buf)
+        op = self.op()
+        title = getattr(op, "name", type(op).__name__)
+
+        table = rich.table.Table(title=f"Summary of {title}\n{n:d} rows")
+
+        table.add_column("Name", justify="left")
+        table.add_column("Type", justify="left")
+        table.add_column("# Nulls", justify="right")
+        table.add_column("% Nulls", justify="right")
+
+        for column, non_nulls in items:
+            table.add_row(
+                column,
+                rich.pretty.Pretty(schema[column]),
+                str(n - non_nulls),
+                f"{100 * (1.0 - non_nulls / n):>3.2f}",
+            )
+
+        with console.capture() as capture:
+            console.print(table)
+        buf.write(capture.get())
 
     def set_column(self, name: str, expr: ir.Value) -> Table:
         """Replace an existing column with a new expression.

poetry.lock

@@ -35,7 +35,7 @@ pandas = ">=1.2.5,<2"
 parsy = ">=1.3.0,<2"
 pydantic = ">=1.9.0,<2"
 regex = ">=2021.7.6"
-tabulate = ">=0.8.9,<1"
+rich = ">=12.4.4,<13"
 toolz = ">=0.11,<0.13"
 clickhouse-cityhash = { version = ">=1.0.2,<2", optional = true }
 clickhouse-driver = { version = ">=0.1,<0.3", optional = true, extras = [