fix: add html2text + Rich render_display_data() for display_data output (#58)

Previously text/html display_data output was either ignored (log) or
printed raw. Now:

- New dependency: html2text (>=2024.2.26)
- New util: render_display_data() in utils.py — returns Rich renderables
  (Markdown or Text) using priority text/markdown > text/html > text/plain.
  text/html is converted via html2text; text/plain is wrapped with
  Text.from_ansi to handle embedded ANSI escapes.
- All three call sites (exec.py, automation.py, repl.py) refactored to
  use the shared function instead of duplicated if/elif chains.
- Each module uses a shared Console instance (_console / self.console)
  to avoid re-allocating per-output during streaming.
- Removed direct html2text and Markdown imports from the three call sites
  — they now just pass the renderable to Console.print().
- Tests in test_utils.py cover Markdown return, Text return, priority,
  and no-text fallback.

Author: merlinran

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
     "filelock>=3.29.2",
     "google-auth>=2.49.1",
     "google-auth-oauthlib>=1.3.0",
+    "html2text>=2024.2.26",
     "jupyter-kernel-client",
     "nbformat>=5.10.4",
     "packaging>=24.0",

src/colab_cli/commands/automation.py

@@ -18,12 +18,15 @@
 import json
 from typing import Optional, List
 import typer
+from rich.console import Console
 from typing_extensions import Annotated
 
 from colab_cli.runtime import ColabRuntime
 from colab_cli.contents import ContentsClient
 from colab_cli.auth import get_credentials
-from colab_cli.utils import get_status_code
+from colab_cli.utils import get_status_code, render_display_data
+
+_console = Console()
 
 
 # Default execute() timeout for human-in-the-loop automations (auth /
@@ -35,6 +38,7 @@
 INTERACTIVE_AUTOMATION_TIMEOUT_SEC = 600
 
 
+
 def run_automation(
     name: str,
     op: str,
@@ -153,8 +157,9 @@ def drivefs_hook(deserialize_msg, wsclient):
             if "text" in out:
                 sys.stdout.write(out["text"])
             elif "data" in out:
-                if "text/plain" in out["data"]:
-                    typer.echo(out["data"]["text/plain"])
+                text = render_display_data(out["data"])
+                if text is not None:
+                    _console.print(text)
             elif out.get("output_type") == "error":
                 ename = out.get("ename", "Error")
                 evalue = out.get("evalue", "")

src/colab_cli/commands/execution.py

@@ -20,13 +20,16 @@
 import typer
 import uuid
 from nbformat.v4 import new_output
+from rich.console import Console
 from typing import Optional
 from typing_extensions import Annotated
 
 from colab_cli.runtime import ColabRuntime
-from colab_cli.utils import handle_image, is_terminal_error
+from colab_cli.utils import handle_image, is_terminal_error, render_display_data
 from colab_cli.console import connect_console
 
+_console = Console()
+
 TITLE_REGEX = re.compile(r"^\s*#\s*@title\s+(.*)", re.MULTILINE)
 
 
@@ -72,15 +75,17 @@ def save_output(outputs, cell):
             )
 
 
+
 def display_output(out, output_image=None):
     if out.get("output_type") == "stream":
         stream = sys.stderr if out.get("name") == "stderr" else sys.stdout
         stream.write(out.get("text", ""))
         stream.flush()
     elif "data" in out:
         data = out["data"]
-        if text := data.get("text/plain"):
-            typer.echo(text)
+        text = render_display_data(data)
+        if text is not None:
+            _console.print(text)
         if png := data.get("image/png"):
             handle_image(png, "image/png", target_path=output_image)
         elif jpeg := data.get("image/jpeg"):

src/colab_cli/repl.py

@@ -25,9 +25,8 @@
 from rich.text import Text
 
 from colab_cli.runtime import ColabRuntime
-from colab_cli.utils import handle_image
+from colab_cli.utils import handle_image, render_display_data
 
-console = Console()
 
 
 class ColabREPL:
@@ -43,7 +42,7 @@ def __init__(
         self.history_logger = history_logger
         self.output_image = output_image
         self.kb = KeyBindings()
-        self.console = console
+        self.console = Console()
         self.repl_history: List[dict] = []
 
         @self.kb.add("enter")
@@ -91,14 +90,16 @@ def display_output(self, output: dict):
                     image_displayed = True
                     break
 
-            if "text/plain" in data:
-                text = data["text/plain"]
+            text = render_display_data(data)
+            if text is not None:
                 # Skip generic IPython object reprs if we already showed an image
-                if image_displayed and any(
-                    x in text for x in ["<IPython.core.display.Image", "<Figure size"]
-                ):
-                    return
-                self.console.print(Text.from_ansi(text))
+                if isinstance(text, Text) and image_displayed:
+                    if any(
+                        x in text.plain
+                        for x in ["<IPython.core.display.Image", "<Figure size"]
+                    ):
+                        return
+                self.console.print(text)
         elif output.get("output_type") == "error":
             ename = output.get("ename", "Error")
             evalue = output.get("evalue", "")

src/colab_cli/utils.py

@@ -13,12 +13,15 @@
 # limitations under the License.
 
 import base64
+import html2text
 import logging
 import sys
 import tempfile
 
+from typing import Optional, Union
 
-from typing import Optional
+from rich.markdown import Markdown
+from rich.text import Text
 
 
 def get_status_code(e: Exception) -> Optional[int]:
@@ -66,6 +69,22 @@ def print_kitty(image_bytes: bytes):
         logging.exception("Kitty rendering failed")
 
 
+def render_display_data(data: dict) -> Union[Markdown, Text, None]:
+    """Extract the best text representation from a display_data dict.
+
+    Priority: text/markdown > text/html (via html2text) > text/plain.
+    Returns a Rich renderable (Markdown or Text) or None when no text mime
+    type is present.  Callers can pass the result directly to Console.print().
+    """
+    if "text/markdown" in data:
+        return Markdown(data["text/markdown"])
+    if "text/html" in data:
+        return Markdown(html2text.html2text(data["text/html"]))
+    if "text/plain" in data:
+        return Text.from_ansi(data["text/plain"])
+    return None
+
+
 def handle_image(image_b64: str, mime_type: str = "image/png", target_path: str = None):
     image_bytes = base64.b64decode(image_b64)
     # Print inline using Kitty protocol

tests/test_utils.py

@@ -15,7 +15,11 @@
 import base64
 from unittest.mock import MagicMock, patch
 
-from colab_cli.utils import handle_image, print_kitty
+import pytest
+from rich.markdown import Markdown
+from rich.text import Text
+
+from colab_cli.utils import handle_image, print_kitty, render_display_data
 
 
 @patch("colab_cli.utils.sys.stdout.isatty", return_value=True)
@@ -56,3 +60,28 @@ def test_handle_image(mock_print_kitty, mock_tempfile, capsys):
 
     captured = capsys.readouterr()
     assert "/tmp/fake.png" in captured.out
+
+
+@pytest.mark.parametrize(
+    "data, expected_markup",
+    [
+        ({"text/markdown": "**md**"}, "**md**"),
+        ({"text/html": "<b>hi</b>"}, "**hi**\n\n"),
+        ({"text/markdown": "**md**", "text/html": "<b>hi</b>"}, "**md**"),
+        ({"text/html": "<b>hi</b>", "text/plain": "plain"}, "**hi**\n\n"),
+    ],
+)
+def test_render_display_data_markdown(data, expected_markup):
+    result = render_display_data(data)
+    assert isinstance(result, Markdown)
+    assert result.markup == expected_markup
+
+
+def test_render_display_data_plain():
+    result = render_display_data({"text/plain": "plain"})
+    assert isinstance(result, Text)
+    assert result.plain == "plain"
+
+
+def test_render_display_data_none():
+    assert render_display_data({"image/png": "..."}) is None