feat: add max_columns option for table repr

Author: jcrist

ibis/backends/tests/test_client.py

@@ -9,6 +9,7 @@
 import pandas as pd
 import pandas.testing as tm
 import pytest
+import rich.console
 import sqlalchemy as sa
 from pytest import mark, param
 
@@ -839,64 +840,94 @@ def test_dunder_array_column(alltypes, dtype):
 
 
 @pytest.mark.parametrize("interactive", [True, False])
-def test_repr(alltypes, interactive):
+def test_repr(alltypes, interactive, monkeypatch):
+    monkeypatch.setattr(ibis.options, "interactive", interactive)
+
     expr = alltypes.select("id", "int_col")
 
     val = str(alltypes.limit(5).id.execute().iloc[0])
 
-    old = ibis.options.interactive
-    ibis.options.interactive = interactive
-    try:
-        s = repr(expr)
-        # no control characters
-        assert all(c.isprintable() or c in "\n\r\t" for c in s)
-        assert "id" in s
-        if interactive:
-            assert val in s
-        else:
-            assert val not in s
-    finally:
-        ibis.options.interactive = old
+    s = repr(expr)
+    # no control characters
+    assert all(c.isprintable() or c in "\n\r\t" for c in s)
+    assert "id" in s
+    if interactive:
+        assert val in s
+    else:
+        assert val not in s
 
 
 @pytest.mark.parametrize("show_types", [True, False])
-def test_interactive_repr_show_types(alltypes, show_types):
+def test_interactive_repr_show_types(alltypes, show_types, monkeypatch):
+    monkeypatch.setattr(ibis.options, "interactive", True)
+    monkeypatch.setattr(ibis.options.repr.interactive, "show_types", show_types)
+
     expr = alltypes.select("id")
-    old = ibis.options.interactive
-    ibis.options.interactive = True
-    ibis.options.repr.interactive.show_types = show_types
-    try:
-        s = repr(expr)
-        if show_types:
-            assert "int" in s
-        else:
-            assert "int" not in s
-    finally:
-        ibis.options.interactive = old
+    s = repr(expr)
+    if show_types:
+        assert "int" in s
+    else:
+        assert "int" not in s
+
+
+@pytest.mark.parametrize("is_jupyter", [True, False])
+def test_interactive_repr_max_columns(alltypes, is_jupyter, monkeypatch):
+    monkeypatch.setattr(ibis.options, "interactive", True)
+
+    cols = {f"c_{i}": ibis._.id + i for i in range(50)}
+    expr = alltypes.mutate(**cols).select(*cols)
+
+    console = rich.console.Console(force_jupyter=is_jupyter)
+    console.options.max_width = 80
+    options = console.options.copy()
+
+    # max_columns = 0
+    text = "".join(s.text for s in console.render(expr, options))
+    assert " c_0 " in text
+    if is_jupyter:
+        # Default of 20 columns are written
+        assert " c_19 " in text
+    else:
+        # width calculations truncates well before 20 columns
+        assert " c_19 " not in text
+
+    # max_columns = 3
+    monkeypatch.setattr(ibis.options.repr.interactive, "max_columns", 3)
+    text = "".join(s.text for s in console.render(expr, options))
+    assert " c_2 " in text
+    assert " c_3 " not in text
+
+    # max_columns = None
+    monkeypatch.setattr(ibis.options.repr.interactive, "max_columns", None)
+    text = "".join(s.text for s in console.render(expr, options))
+    assert " c_0 " in text
+    if is_jupyter:
+        # All columns written
+        assert " c_49 " in text
+    else:
+        # width calculations still truncates
+        assert " c_19 " not in text
 
 
 @pytest.mark.parametrize("expr_type", ["table", "column"])
 @pytest.mark.parametrize("interactive", [True, False])
-def test_repr_mimebundle(alltypes, interactive, expr_type):
+def test_repr_mimebundle(alltypes, interactive, expr_type, monkeypatch):
+    monkeypatch.setattr(ibis.options, "interactive", interactive)
+
     if expr_type == "column":
         expr = alltypes.id
     else:
         expr = alltypes.select("id", "int_col")
 
     val = str(alltypes.limit(5).id.execute().iloc[0])
 
-    old = ibis.options.interactive
-    ibis.options.interactive = interactive
-    try:
-        reprs = expr._repr_mimebundle_(include=["text/plain", "text/html"], exclude=[])
-        for format in ["text/plain", "text/html"]:
-            assert "id" in reprs[format]
-            if interactive:
-                assert val in reprs[format]
-            else:
-                assert val not in reprs[format]
-    finally:
-        ibis.options.interactive = old
+    reprs = expr._repr_mimebundle_(include=["text/plain", "text/html"], exclude=[])
+    for format in ["text/plain", "text/html"]:
+        assert "id" in reprs[format]
+        if interactive:
+            assert val in reprs[format]
+        else:
+            assert val not in reprs[format]
 
 
 @pytest.mark.never(

ibis/config.py

@@ -77,6 +77,10 @@ class Interactive(Config):
     ----------
     max_rows : int
         Maximum rows to pretty print.
+    max_columns : int | None
+        The maximum number of columns to pretty print. If 0 (the default), the
+        number of columns will be inferred from output console size. Set to
+        `None` for no limit.
     max_length : int
         Maximum length for pretty-printed arrays and maps.
     max_string : int
@@ -88,6 +92,7 @@ class Interactive(Config):
     """
 
     max_rows: int = 10
+    max_columns: Optional[int] = 0
     max_length: int = 2
     max_string: int = 80
     max_depth: int = 1

ibis/expr/types/pretty.py

@@ -197,10 +197,18 @@ def to_rich_table(table, console_width=None):
 
     orig_ncols = len(table.columns)
 
-    # First determine the maximum subset of columns that *might* fit in the
-    # current console. Note that not every column here may actually fit later
-    # on once we know the repr'd width of the data.
-    if console_width:
+    max_columns = ibis.options.repr.interactive.max_columns
+    if console_width == float("inf"):
+        if max_columns == 0:
+            # Show up to 20 columns by default, mirroring pandas's behavior
+            if orig_ncols >= 20:
+                table = table.select(*table.columns[:20])
+        elif max_columns is not None and max_columns < orig_ncols:
+            table = table.select(*table.columns[:max_columns])
+    else:
+        # Determine the maximum subset of columns that *might* fit in the
+        # current console. Note that not every column here may actually fit
+        # later on once we know the repr'd width of the data.
         computed_cols = []
         remaining = console_width - 1  # 1 char for left boundary
         for c in table.columns:
@@ -211,8 +219,12 @@ def to_rich_table(table, console_width=None):
                 computed_cols.append(c)
                 remaining -= needed
             else:
-                table = table.select(*computed_cols)
                 break
+        if max_columns not in (0, None):
+            # If an explicit limit on max columns is set, apply it
+            computed_cols = computed_cols[:max_columns]
+        if orig_ncols > len(computed_cols):
+            table = table.select(*computed_cols)
 
     # Compute the data and return a pandas dataframe
     nrows = ibis.options.repr.interactive.max_rows
@@ -269,29 +281,35 @@ def to_rich_table(table, console_width=None):
     # figure it out for us.
     columns_truncated = orig_ncols > len(col_info)
     col_widths = {}
-    flex_cols = []
-    remaining = console_width - 1
-    if columns_truncated:
-        remaining -= 4
-    for name, _, min_width, max_width in col_info:
-        remaining -= min_width + 3
-        col_widths[name] = min_width
-        if min_width != max_width:
-            flex_cols.append((name, max_width))
-
-    while True:
-        next_flex_cols = []
-        for name, max_width in flex_cols:
-            if remaining:
-                remaining -= 1
-                if max_width is not None:
-                    col_widths[name] += 1
-                if max_width is None or col_widths[name] < max_width:
-                    next_flex_cols.append((name, max_width))
-            else:
+    if console_width == float("inf"):
+        # Always use the max_width if there's infinite console space
+        for name, _, _, max_width in col_info:
+            col_widths[name] = max_width
+    else:
+        # Allocate the remaining space evenly between the flexible columns
+        flex_cols = []
+        remaining = console_width - 1
+        if columns_truncated:
+            remaining -= 4
+        for name, _, min_width, max_width in col_info:
+            remaining -= min_width + 3
+            col_widths[name] = min_width
+            if min_width != max_width:
+                flex_cols.append((name, max_width))
+
+        while True:
+            next_flex_cols = []
+            for name, max_width in flex_cols:
+                if remaining:
+                    remaining -= 1
+                    if max_width is not None:
+                        col_widths[name] += 1
+                    if max_width is None or col_widths[name] < max_width:
+                        next_flex_cols.append((name, max_width))
+                else:
+                    break
+            if not next_flex_cols:
                 break
-        if not next_flex_cols:
-            break
 
     rich_table = rich.table.Table(padding=(0, 1, 0, 1))
 
@@ -313,6 +331,7 @@ def to_rich_table(table, console_width=None):
             justify="left",
             vertical="middle",
             width=1,
+            min_width=1,
             no_wrap=True,
         )
 

ibis/expr/types/relations.py

@@ -126,7 +126,19 @@ def __rich_console__(self, console, options):
         if not ibis.options.interactive:
             return console.render(Text(self._repr()), options=options)
 
-        table = to_rich_table(self, options.max_width)
+        if console.is_jupyter:
+            # Rich infers a console width in jupyter notebooks, but since
+            # notebooks can use horizontal scroll bars we don't want to apply a
+            # limit here. Since rich requires an integer for max_width, we
+            # choose an arbitrarily large integer bound. Note that we need to
+            # handle this here rather than in `to_rich_table`, as this setting
+            # also needs to be forwarded to `console.render`.
+            options = options.update(max_width=1_000_000)
+            width = None
+        else:
+            width = options.max_width
+
+        table = to_rich_table(self, width)
         return console.render(table, options=options)
 
     def __getitem__(self, what):