Skip to content

Commit

Permalink
feat(ux): use rich to format Table.info() output
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Aug 2, 2022
1 parent 10a659d commit 67234c3
Show file tree
Hide file tree
Showing 7 changed files with 61 additions and 34 deletions.
22 changes: 16 additions & 6 deletions ibis/backends/tests/test_generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import importlib
import io
import operator
from contextlib import redirect_stdout

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -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(
Expand Down
4 changes: 4 additions & 0 deletions ibis/common/grounds.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):

Expand Down
45 changes: 28 additions & 17 deletions ibis/expr/types/relations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down
12 changes: 7 additions & 5 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
8 changes: 4 additions & 4 deletions requirements.txt

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion setup.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 67234c3

Please sign in to comment.