diff --git a/src/datalab_api/_base.py b/src/datalab_api/_base.py index 7275a92..5a9c813 100644 --- a/src/datalab_api/_base.py +++ b/src/datalab_api/_base.py @@ -7,8 +7,10 @@ from typing import Any, Optional import httpx +from rich.console import Console from rich.logging import RichHandler from rich.pretty import pprint +from rich.table import Table __version__ = version("datalab-api") @@ -16,17 +18,46 @@ def pretty_displayer(method): + """A decorator which wraps a method with a 'display' kwarg, which will + either pretty print a JSON response, or display a Rich table. + + """ + @functools.wraps(method) def rich_wrapper(self, *args, **kwargs): display = kwargs.pop("display", False) + page_limit = kwargs.pop("page_limit", 10) result = method(self, *args, **kwargs) if display: - if isinstance(result, dict) and "blocks_obj" in result: - blocks: dict[str, dict] = result["blocks_obj"] - for block in blocks.values(): - if "bokeh_plot_data" in block: - bokeh_from_json(block) - pprint(result, max_length=None, max_string=100, max_depth=3) + blocks = None + if isinstance(result, dict): + if "blocks_obj" in result: + blocks = result["blocks_obj"] + elif "blocktype" in result: + blocks = {result["block_id"]: result} + result = None + if blocks: + for block in blocks.values(): + if "bokeh_plot_data" in block: + bokeh_from_json(block) + if result: + pprint(result, max_length=None, max_string=100, max_depth=3) + elif isinstance(result, list): + table = Table(show_lines=True) + table.add_column("type", overflow="crop", justify="center") + table.add_column("ID", style="cyan", no_wrap=True) + table.add_column("refcode", style="magenta", no_wrap=True) + table.add_column("name", width=30, overflow="ellipsis", no_wrap=True) + for item in result[:page_limit]: + table.add_row( + item["type"][0].upper(), + item["item_id"], + item["refcode"], + item["name"], + end_section=True, + ) + console = Console() + console.print(table) return result @@ -41,16 +72,19 @@ def __new__(cls, name, bases, dct): return super().__new__(cls, name, bases, dct) -def bokeh_from_json(block_data): +def bokeh_from_json(block_data, show=True): from bokeh.io import curdoc - from bokeh.plotting import show + from bokeh.plotting import show as bokeh_show if "bokeh_plot_data" in block_data: bokeh_plot_data = block_data["bokeh_plot_data"] else: bokeh_plot_data = block_data curdoc().replace_with_json(bokeh_plot_data["doc"]) - show(curdoc().roots[0]) + if show: + bokeh_show(curdoc().roots[0]) + + return curdoc() class BaseDatalabClient(metaclass=AutoPrettyPrint):