diff --git a/ibis/expr/types/core.py b/ibis/expr/types/core.py index 043f8385474df..7e37f2b6113e1 100644 --- a/ibis/expr/types/core.py +++ b/ibis/expr/types/core.py @@ -77,20 +77,9 @@ def __repr__(self) -> str: if not opts.interactive: return self._repr() - from ibis.expr.types.pretty import simple_console - - with simple_console.capture() as capture: - try: - simple_console.print(self) - except TranslationError as e: - lines = [ - "Translation to backend failed", - f"Error message: {e.args[0]}", - "Expression repr follows:", - self._repr(), - ] - return "\n".join(lines) - return capture.get().rstrip() + from ibis.expr.types.pretty import pretty_repr + + return pretty_repr(self) def __reduce__(self): return (self.__class__, (self._arg,)) diff --git a/ibis/expr/types/generic.py b/ibis/expr/types/generic.py index 03af02ef5ff42..18f36950eed46 100644 --- a/ibis/expr/types/generic.py +++ b/ibis/expr/types/generic.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Literal from public import public +from rich.table import Table as RichTable import ibis import ibis.common.exceptions as com @@ -1293,9 +1294,69 @@ def __array__(self, dtype=None): return self.execute().__array__(dtype) def __interactive_rich_console__(self, console, options): + return console.render(self.preview(), options=options) + + def preview( + self, + *, + max_rows: int | None = None, + show_types: bool | None = None, + max_length: int | None = None, + max_string: int | None = None, + max_depth: int | None = None, + console_width: int | float | None = None, + ) -> RichTable: + """Pretty-print a preview using rich. + + This is an explicit version of what you get when you inspect + this object in interactive mode, except with this version you + can pass formatting options. The options are the same as those exposed + in `ibis.options.interactive`. + + Parameters + ---------- + max_rows + Maximum number of rows to display + show_types + Whether to show column types in the 2nd row, under the column name + max_length + Maximum length for pretty-printed arrays and maps. + max_string + Maximum length for pretty-printed strings. + max_depth + Maximum depth for nested data types. + console_width + Width of the console in characters. If not specified, the width + will be inferred from the console. + + Examples + -------- + >>> import ibis + >>> ibis.options.interactive = False + >>> t = ibis.examples.penguins.fetch() + >>> t.island.preview(max_rows=3, show_types=False, max_string=5) + ┏━━━━━━━━┓ + ┃ island ┃ + ┡━━━━━━━━┩ + │ Torg… │ + │ Torg… │ + │ Torg… │ + │ … │ + └────────┘ + """ + from ibis.expr.types.pretty import to_rich_table + named = self.name(self.op().name) projection = named.as_table() - return console.render(projection, options=options) + return to_rich_table( + projection, + max_rows=max_rows, + show_types=show_types, + max_length=max_length, + max_string=max_string, + max_depth=max_depth, + console_width=console_width, + ) def __pyarrow_result__( self, table: pa.Table, data_mapper: type[PyArrowData] | None = None diff --git a/ibis/expr/types/pretty.py b/ibis/expr/types/pretty.py index 31354895a0be2..32b547b7cf7ed 100644 --- a/ibis/expr/types/pretty.py +++ b/ibis/expr/types/pretty.py @@ -10,35 +10,65 @@ from rich.align import Align from rich.console import Console from rich.text import Text +from rich.table import Table as RichTable import ibis import ibis.expr.datatypes as dt +from ibis.expr.types import Table +from ibis.common.exceptions import TranslationError -# A console with all color/markup disabled, used for `__repr__` -simple_console = Console(force_terminal=False) - -def _format_nested(values): - interactive = ibis.options.repr.interactive +def pretty_repr(expr: ibis.Expr) -> str: + # All color/markup disabled, used for `__repr__` + console = Console(force_terminal=False) + with console.capture() as capture: + try: + console.print(expr) + except TranslationError as e: + lines = [ + "Translation to backend failed", + f"Error message: {e.args[0]}", + "Expression repr follows:", + expr._repr(), + ] + return "\n".join(lines) + return capture.get().rstrip() + + +def _format_nested( + values, + *, + max_length: int | None = None, + max_string: int | None = None, + max_depth: int | None = None, +): + if max_length is None: + max_length = ibis.options.repr.interactive.max_length + if max_string is None: + max_string = ibis.options.repr.interactive.max_string + if max_depth is None: + max_depth = ibis.options.repr.interactive.max_depth return [ rich.pretty.Pretty( v, - max_length=interactive.max_length, - max_string=interactive.max_string, - max_depth=interactive.max_depth, + max_length=max_length, + max_string=max_string, + max_depth=max_depth, ) for v in values ] @singledispatch -def format_values(dtype, values): - return _format_nested(values) +def format_values(dtype, values, **fmt_kwargs): + return _format_nested(values, **fmt_kwargs) @format_values.register(dt.Map) -def _(dtype, values): - return _format_nested([None if v is None else dict(v) for v in values]) +def _(dtype, values, **fmt_kwargs): + return _format_nested( + [None if v is None else dict(v) for v in values], **fmt_kwargs + ) @format_values.register(dt.GeoSpatial) @@ -49,7 +79,7 @@ def _(dtype, values): @format_values.register(dt.JSON) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): def try_json(v): if v is None: return None @@ -58,17 +88,17 @@ def try_json(v): except Exception: # noqa: BLE001 return v - return _format_nested([try_json(v) for v in values]) + return _format_nested([try_json(v) for v in values], **fmt_kwargs) @format_values.register(dt.Boolean) @format_values.register(dt.UUID) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): return [Text(str(v)) for v in values] @format_values.register(dt.Decimal) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): if dtype.scale is not None: fmt = f"{{:.{dtype.scale}f}}" return [Text.styled(fmt.format(v), "bold cyan") for v in values] @@ -78,12 +108,12 @@ def _(dtype, values): @format_values.register(dt.Integer) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): return [Text.styled(str(int(v)), "bold cyan") for v in values] @format_values.register(dt.Floating) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): floats = [float(v) for v in values] # Extract and format all finite floats finites = [f for f in floats if isfinite(f)] @@ -102,7 +132,7 @@ def _(dtype, values): @format_values.register(dt.Timestamp) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): if all(v.microsecond == 0 for v in values): timespec = "seconds" elif all(v.microsecond % 1000 == 0 for v in values): @@ -115,13 +145,13 @@ def _(dtype, values): @format_values.register(dt.Date) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): dates = [v.date() if isinstance(v, datetime.datetime) else v for v in values] return [Text.styled(d.isoformat(), "magenta") for d in dates] @format_values.register(dt.Time) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): times = [v.time() if isinstance(v, datetime.datetime) else v for v in values] if all(t.microsecond == 0 for t in times): timespec = "seconds" @@ -133,7 +163,7 @@ def _(dtype, values): @format_values.register(dt.Interval) -def _(dtype, values): +def _(dtype, values, **fmt_kwargs): return [Text.styled(str(v), "magenta") for v in values] @@ -149,8 +179,9 @@ def _(dtype, values): @format_values.register(dt.String) -def _(dtype, values): - max_string = ibis.options.repr.interactive.max_string +def _(dtype, values, *, max_string: int | None = None, **fmt_kwargs): + if max_string is None: + max_string = ibis.options.repr.interactive.max_string out = [] for v in values: v = str(v) @@ -181,7 +212,14 @@ def _(dtype, values): return out -def format_column(dtype, values): +def format_column( + dtype, + values, + *, + max_length: int | None = None, + max_string: int | None = None, + max_depth: int | None = None, +): import pandas as pd null_str = Text.styled("NULL", style="dim") @@ -199,7 +237,13 @@ def isnull(x): nonnull = [v for v in values if not isnull(v)] if nonnull: - formatted = format_values(dtype, nonnull) + formatted = format_values( + dtype, + nonnull, + max_length=max_length, + max_string=max_string, + max_depth=max_depth, + ) next_f = iter(formatted).__next__ out = [null_str if isnull(v) else next_f() for v in values] else: @@ -227,13 +271,28 @@ def format_dtype(dtype): return Text.styled(strtyp, "dim") -def to_rich_table(table, console_width=None): +def to_rich_table( + table: Table, + *, + max_rows: int | None = None, + max_columns: int | None = None, + show_types: bool | None = None, + max_length: int | None = None, + max_string: int | None = None, + max_depth: int | None = None, + console_width: int | float | None = None, +) -> RichTable: + if max_rows is None: + max_rows = ibis.options.repr.interactive.max_rows + if max_columns is None: + max_columns = ibis.options.repr.interactive.max_columns + if show_types is None: + show_types = ibis.options.repr.interactive.show_types if console_width is None: console_width = float("inf") orig_ncols = len(table.columns) - max_columns = ibis.options.repr.interactive.max_columns if console_width == float("inf"): # When there's infinite display space, only subset columns # if an explicit limit has been set. @@ -260,10 +319,7 @@ def to_rich_table(table, console_width=None): 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 - result = table.limit(nrows + 1).to_pyarrow() - + result = table.limit(max_rows + 1).to_pyarrow() # Now format the columns in order, stopping if the console width would # be exceeded. col_info = [] @@ -272,12 +328,14 @@ def to_rich_table(table, console_width=None): remaining = console_width - 1 # 1 char for left boundary for name, dtype in table.schema().items(): formatted, min_width, max_width = format_column( - dtype, result[name].to_pylist()[:nrows] + dtype, + result[name].to_pylist()[:max_rows], + max_length=max_length, + max_string=max_string, + max_depth=max_depth, ) dtype_str = format_dtype(dtype) - if ibis.options.repr.interactive.show_types and not isinstance( - dtype, (dt.Struct, dt.Map, dt.Array) - ): + if show_types and not isinstance(dtype, (dt.Struct, dt.Map, dt.Array)): # Don't truncate non-nested dtypes min_width = max(min_width, len(dtype_str)) @@ -375,7 +433,7 @@ def add_row(*args, **kwargs): else: add_row = rich_table.add_row - if ibis.options.repr.interactive.show_types: + if show_types: add_row( *(Align(s, align="left") for s in formatted_dtypes), end_section=True, @@ -385,7 +443,7 @@ def add_row(*args, **kwargs): add_row(*row) # If the rows are truncated, add a trailing ellipsis row - if len(result) > nrows: + if len(result) > max_rows: rich_table.add_row( *(Align("[dim]…[/]", align=c.justify) for c in rich_table.columns) ) diff --git a/ibis/expr/types/relations.py b/ibis/expr/types/relations.py index c7d2b47bed656..265d2b34f9e52 100644 --- a/ibis/expr/types/relations.py +++ b/ibis/expr/types/relations.py @@ -12,6 +12,7 @@ import toolz from public import public +from rich.table import Table as RichTable import ibis import ibis.common.exceptions as com @@ -352,8 +353,6 @@ def _cast(self, schema: SupportsSchema, cast_method: str = "cast") -> Table: return self.select(*cols) def __interactive_rich_console__(self, console, options): - from ibis.expr.types.pretty import to_rich_table - 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 @@ -367,7 +366,7 @@ def __interactive_rich_console__(self, console, options): width = options.max_width try: - table = to_rich_table(self, width) + table = self.preview(console_width=width) except Exception as e: # In IPython exceptions inside of _repr_mimebundle_ are swallowed to # allow calling several display functions and choosing to display