From a3aa23684346b88aa05d76d0a712f98310ed4eaa Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Wed, 8 Feb 2023 16:38:44 -0600 Subject: [PATCH] feat: add `max_columns` option for table repr --- ibis/backends/tests/test_client.py | 109 ++++++++++++++++++----------- ibis/config.py | 5 ++ ibis/expr/types/pretty.py | 73 ++++++++++++------- ibis/expr/types/relations.py | 14 +++- 4 files changed, 134 insertions(+), 67 deletions(-) diff --git a/ibis/backends/tests/test_client.py b/ibis/backends/tests/test_client.py index c3b7e361ea9c..1d7b09d6e2d4 100644 --- a/ibis/backends/tests/test_client.py +++ b/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,45 +840,80 @@ 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: @@ -885,18 +921,13 @@ def test_repr_mimebundle(alltypes, interactive, expr_type): 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( diff --git a/ibis/config.py b/ibis/config.py index 379c8f9bef01..65dbb666e7f9 100644 --- a/ibis/config.py +++ b/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 diff --git a/ibis/expr/types/pretty.py b/ibis/expr/types/pretty.py index d39c97163453..fdc5ba1c4b1e 100644 --- a/ibis/expr/types/pretty.py +++ b/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, ) diff --git a/ibis/expr/types/relations.py b/ibis/expr/types/relations.py index d5dec5b77d77..339ec45563d2 100644 --- a/ibis/expr/types/relations.py +++ b/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):