Add ability to show rich tables outside of CLI

src/datalab_api/_base.py

@@ -7,26 +7,57 @@
 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")
 
 __all__ = ("__version__", "BaseDatalabClient")
 
 
 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):