feat(ux): improve rich rendering (#11326)

Co-authored-by: Phillip Cloud <[email protected]>

Author: NickCrews

ibis/backends/tests/test_client.py

@@ -14,6 +14,7 @@
 from typing import TYPE_CHECKING
 
 import pytest
+import rich
 import rich.console
 import sqlglot as sg
 import toolz
@@ -1112,21 +1113,84 @@ def test_dunder_array_column(alltypes, dtype):
     np.testing.assert_array_equal(result, expected)
 
 
-@pytest.mark.parametrize("interactive", [True, False])
-def test_repr(alltypes, interactive, monkeypatch):
-    pytest.importorskip("rich")
-
+@pytest.mark.parametrize(
+    "interactive", [True, False], ids=["interactive", "non_interactive"]
+)
+@pytest.mark.parametrize("is_jupyter", [True, False], ids=["jupyter", "not_jupyter"])
+@pytest.mark.parametrize(
+    ("method_name", "method"),
+    [
+        ("repr", repr),
+        (
+            "mimebundle",
+            lambda wide_table: wide_table._repr_mimebundle_(["text/plain"], [])[
+                "text/plain"
+            ],
+        ),
+        (
+            "preview",
+            lambda wide_table: wide_table.preview()._repr_mimebundle_(None, None)[
+                "text/plain"
+            ],
+        ),
+    ],
+    ids=["repr", "mimebundle", "preview"],
+)
+def test_reprs(alltypes, interactive, is_jupyter, method_name, method, monkeypatch):
     monkeypatch.setattr(ibis.options, "interactive", interactive)
 
-    expr = alltypes.select("date_string_col")
+    # Depending on the order that tests are run, someone may have already
+    # called rich.get_console() and created the default console,
+    # which will have inferred is_jupyter from the environment,
+    # leading to False in pytest.
+    # 1. Make it so any newly-created Consoles will use the right value of `is_jupyter`
+    # 2. Update the default console to use the right value of `is_jupyter`
+    monkeypatch.setattr("rich.console._is_jupyter", lambda: is_jupyter)
+    default_console = rich.get_console()
+    new_console = rich.console.Console(force_jupyter=is_jupyter)
+    monkeypatch.setattr(default_console, "__dict__", new_console.__dict__)
+    expected_color_system = ("truecolor",) if is_jupyter else (None, "standard")
+    assert new_console.is_jupyter == is_jupyter
+    assert default_console.is_jupyter == is_jupyter
+    assert new_console.color_system in expected_color_system
+    assert default_console.color_system in expected_color_system
+
+    wide_table = alltypes.mutate(
+        *[alltypes[c].name(f"{c}_2") for c in alltypes.columns],
+        *[alltypes[c].name(f"{c}_3") for c in alltypes.columns],
+    )
 
-    s = repr(expr)
-    # no control characters
-    assert all(c.isprintable() or c in "\n\r\t" for c in s)
-    if interactive:
-        assert "/" in s
+    s = method(wide_table)
+
+    color_or_control_codes = {c for c in s if not c.isprintable() and c not in "\n\r\t"}
+    if (method_name == "mimebundle" and interactive and is_jupyter) or (
+        method_name == "preview" and is_jupyter
+    ):
+        assert color_or_control_codes
     else:
-        assert "/" not in s
+        assert not color_or_control_codes
+
+    # ┃ characters separate columns in the table
+    n_columns_seen = max(line.count("┃") for line in s.splitlines()) - 1
+    n_columns_seen = max(n_columns_seen, 0)
+
+    n_wide_columns = len(wide_table.columns)
+    n_original_columns = len(alltypes.columns)
+
+    assert (
+        # preview always shows exactly what was requested
+        (method_name == "preview" and n_columns_seen == n_wide_columns)
+        or (
+            interactive
+            and (
+                # table should have unbounded width, every column should be shown
+                (is_jupyter and n_columns_seen == n_wide_columns)
+                # table should be truncated to fit the console width
+                or n_columns_seen < n_original_columns
+            )
+        )
+        or not n_columns_seen
+    )
 
 
 @pytest.mark.parametrize("show_types", [True, False])

ibis/expr/types/_rich.py

@@ -412,7 +412,28 @@ def to_rich_table(
             if not next_flex_cols:
                 break
 
-    rich_table = rich.table.Table(padding=(0, 1, 0, 1))
+    # If the console width is unbounded, then explicitly make each column
+    # take up its full width.
+    # Otherwise, when you render the returned rich.Table in a rich.Console,
+    # if the Console width is skinny, each table might be scrunched.
+    # This happens for example when display()ing a rich.Table in a Jupyter notebook,
+    # and you end up with
+    # ┏┳┳━━━━━┳━━━━┳━━━━━━━━┳━━┳┳┳┳┳┳┳┳┳┳┳┳┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳━┳┳━┳┳━┳┳━┓
+    # ┃┃┃ bi… ┃ b… ┃ flipp… ┃  ┃┃┃┃┃┃┃┃┃┃┃┃┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃┃ ┃┃ ┃┃ ┃
+    # ┡╇╇━━━━━╇━━━━╇━━━━━━━━╇━━╇╇╇╇╇╇╇╇╇╇╇╇╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇━╇╇━╇╇━╇╇━┩
+    # │││ fl… │ f… │ int64  │  │││││││││││││ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ││ ││ ││ │
+    # ├┼┼─────┼────┼────────┼──┼┼┼┼┼┼┼┼┼┼┼┼┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼─┼┼─┼┼─┼┼─┤
+    # │││ 39… │ 1… │    181 │  │││││││││││││ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ││ ││ ││ │
+    # │││ 39… │ 1… │    186 │  │││││││││││││ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ││ ││ ││ │
+    # │││ 40… │ 1… │    195 │  │││││││││││││ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ││ ││ ││ │
+    # └┴┴─────┴────┴────────┴──┴┴┴┴┴┴┴┴┴┴┴┴┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴─┴┴─┴┴─┴┴─┘
+    # unless you set the width of the table explicitly.
+    table_width = (
+        1 + sum(col_widths.values()) + len(col_info) * 3
+        if console_width == float("inf")
+        else None
+    )
+    rich_table = rich.table.Table(padding=(0, 1, 0, 1), width=table_width)
 
     # Configure the columns on the rich table.
     for name, dtype, _, max_width in col_info:

ibis/expr/types/core.py

@@ -51,7 +51,7 @@ def _noninteractive_repr(self) -> str:
 
     def __repr__(self) -> str:
         if ibis.options.interactive:
-            return capture_rich_renderable(self)
+            return capture_rich_renderable(self, no_color=True)["text/plain"]
         else:
             return self._noninteractive_repr()
 

ibis/expr/types/generic.py

@@ -14,7 +14,7 @@
 from ibis.common.grounds import Singleton
 from ibis.expr.rewrites import rewrite_window_input
 from ibis.expr.types.core import Expr, _binop
-from ibis.expr.types.rich import FixedTextJupyterMixin, to_rich
+from ibis.expr.types.rich import RichJupyterMixin, to_rich
 from ibis.util import deprecated, experimental, promote_list
 
 if TYPE_CHECKING:
@@ -1537,7 +1537,7 @@ def _repr_html_(self) -> str | None:
 
 
 @public
-class Column(Value, FixedTextJupyterMixin):
+class Column(Value, RichJupyterMixin):
     # Higher than numpy objects
     __array_priority__ = 20
 

ibis/expr/types/relations.py

@@ -23,7 +23,7 @@
 from ibis.expr.rewrites import DerefMap
 from ibis.expr.types.core import Expr
 from ibis.expr.types.generic import Value, literal
-from ibis.expr.types.rich import FixedTextJupyterMixin, to_rich
+from ibis.expr.types.rich import RichJupyterMixin, to_rich
 from ibis.expr.types.temporal import TimestampColumn
 from ibis.util import deprecated, experimental
 
@@ -144,7 +144,7 @@ def unwrap_aliases(values: Iterator[ir.Value]) -> Mapping[str, ir.Value]:
 
 
 @public
-class Table(Expr, FixedTextJupyterMixin):
+class Table(Expr, RichJupyterMixin):
     """An immutable and lazy dataframe.
 
     Analogous to a SQL table or a pandas DataFrame. A table expression contains

ibis/expr/types/rich.py

@@ -6,45 +6,52 @@
 
 from __future__ import annotations
 
+import contextlib
 from typing import TYPE_CHECKING
 
 from ibis.expr import types as ir
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     import rich.panel
     import rich.table
     from rich.console import RenderableType
 
 
 try:
-    from rich.jupyter import JupyterMixin
+    import rich
 except ImportError:
 
-    class FixedTextJupyterMixin:
+    class RichJupyterMixin:
         """No-op when rich is not installed."""
 else:
 
-    class FixedTextJupyterMixin(JupyterMixin):
-        """JupyterMixin adds a spurious newline to text, this fixes the issue."""
+    class RichJupyterMixin:
+        """Adds `_repr_mimebundle_()` to anything with a `__rich_console__()`."""
 
-        def _repr_mimebundle_(self, *args, **kwargs):
-            try:
-                bundle = super()._repr_mimebundle_(*args, **kwargs)
-            except Exception:  # noqa: BLE001
-                return None
-            else:
-                bundle["text/plain"] = bundle["text/plain"].rstrip()
-                return bundle
+        def _repr_mimebundle_(
+            self, include: Sequence[str], exclude: Sequence[str], **kwargs
+        ) -> dict[str, str]:
+            bundle = capture_rich_renderable(self, no_color=False)
+            return {k: bundle[k] for k in (bundle.keys() & include).difference(exclude)}
 
 
-def capture_rich_renderable(renderable: RenderableType) -> str:
-    """Convert a rich renderable (has a __rich_console__(), etc) to a string."""
+def capture_rich_renderable(
+    renderable: RenderableType, *, no_color: bool
+) -> dict[str, str]:
+    """Convert a rich renderable (has a __rich_console__(), etc) to text and html representations."""
     from rich.console import Console
 
-    console = Console(force_terminal=False)
-    with console.capture() as capture:
-        console.print(renderable)
-    return capture.get().rstrip()
+    color_system = None if no_color else "auto"
+    console = Console()
+    width = console.width
+    if console.is_jupyter:
+        width = 1_000_000
+    with _with_rich_configured(
+        width=width, color_system=color_system, force_terminal=False
+    ):
+        return _RichMimeBundler(renderable).get_mimebundle()
 
 
 def to_rich(
@@ -74,3 +81,67 @@ def to_rich(
             max_depth=max_depth,
             console_width=console_width,
         )
+
+
+@contextlib.contextmanager
+def _with_rich_configured(**config):
+    """Context manager to temporarily configure rich."""
+    from rich.console import Console
+
+    global_console = rich.get_console()
+    new_console = Console(**config)
+    original_config = global_console.__dict__
+
+    try:
+        global_console.__dict__ = new_console.__dict__
+        yield
+    finally:
+        global_console.__dict__ = original_config
+
+
+try:
+    from rich.jupyter import JupyterMixin
+except ImportError:
+    JupyterMixin = object
+
+
+class _RichMimeBundler(JupyterMixin):
+    def __init__(self, renderable: RenderableType):
+        self.renderable = renderable
+
+    def __rich_console__(self, console, options):
+        yield self.renderable
+
+    def get_mimebundle(self) -> dict[str, str]:
+        with _with_rich_display_disabled():
+            bundle = super()._repr_mimebundle_(include=None, exclude=None)
+        bundle["text/plain"] = bundle["text/plain"].rstrip()  # Remove trailing newline
+        return bundle
+
+
+@contextlib.contextmanager
+def _with_rich_display_disabled():
+    """Workaround to keep rich from doing spurious display() calls in Jupyter.
+
+    When you display(ibis.Table), without this, an extra output cell is created
+    in the notebook. With this, there is no extra output cell.
+
+    See https://github.com/Textualize/rich/pull/3329
+    """
+    try:
+        from IPython import display as ipython_display
+    except ImportError:
+        # IPython is not installed, so nothing to do
+        yield
+    else:
+
+        def noop_display(*args, **kwargs):
+            pass
+            # ipython_display.display(*args, **kwargs)
+
+        original_display = ipython_display.display
+        try:
+            ipython_display.display = noop_display
+            yield
+        finally:
+            ipython_display.display = original_display