Rich pixels

.github/workflows/tests.yml

@@ -16,7 +16,8 @@ jobs:
       - name: Install dependencies
         run: |
           pip install poetry==1.2.2
-          poetry install --no-interaction --no-ansi
+          poetry config virtualenvs.create false
+          poetry install --no-interaction --no-ansi --all-extras
       - name: QA with flake8, black, isort and mypy
         run: |
           poetry run pre-commit run --all-files

.pre-commit-config.yaml

@@ -37,6 +37,7 @@ repos:
           - GitPython==3.1.26
           - typing-extensions==4.0.1
           - tomli==2.0.1
+          - types-pillow==9.3.0.1
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.4
     hooks:
@@ -54,8 +55,10 @@ repos:
           - cogapp==3.3.0
           - typer-cli==0.0.12
           - typer==0.3.2
-          - rich==10.16.2
+          - rich==12.6.0
           - pydantic==1.9.0
           - GitPython==3.1.26
           - typing-extensions==4.0.1
           - tomli==2.0.1
+          - rich-pixels==2.0.0
+          - pillow==9.3.0

README.md

@@ -39,7 +39,7 @@ The key features include:
 ## Installation
 
 ```
-pip install databooks
+pip install databooks[images]
 ```
 
 ## Usage

databooks/data_models/rich_helpers.py → databooks/data_models/_rich.py

@@ -1,10 +1,20 @@
 """Rich helpers functions for rich renderables in data models."""
+import base64
 from html.parser import HTMLParser
+from io import BytesIO
 from typing import Any, List, Optional, Tuple
 
 from rich import box
 from rich.table import Table
 
+try:
+    from PIL import Image
+    from rich_pixels import Pixels
+
+    _IMG_INSTALLED = True
+except ImportError:
+    _IMG_INSTALLED = False
+
 HtmlAttr = Tuple[str, Optional[str]]
 
 
@@ -72,3 +82,10 @@ def rich(self, **tbl_kwargs: Any) -> Optional[Table]:
         for row in self.rows:
             table.add_row(*row)
         return table
+
+
+def img2rich(_base64: str, max_size: Tuple[int, int]) -> Pixels:
+    """Get rich pixels from base64 image."""
+    with Image.open(BytesIO(base64.b64decode(_base64))) as image:
+        image.thumbnail(max_size, Image.ANTIALIAS)
+        return Pixels.from_image(image)

databooks/data_models/cell.py

@@ -10,8 +10,13 @@
 from rich.syntax import Syntax
 from rich.text import Text
 
+from databooks.data_models._rich import (
+    _IMG_INSTALLED,
+    HtmlTable,
+    RichHtmlTableError,
+    img2rich,
+)
 from databooks.data_models.base import DatabooksBase
-from databooks.data_models.rich_helpers import HtmlTable, RichHtmlTableError
 from databooks.logging import get_logger
 
 logger = get_logger(__file__)
@@ -157,6 +162,9 @@ def _try_parse_html(s: str) -> Optional[ConsoleRenderable]:
 
         mime_func: Dict[str, Callable[[str], Optional[ConsoleRenderable]]] = {
             "text/html": lambda s: _try_parse_html(s),
+            "image/png": lambda s: img2rich(s, self.max_size)
+            if _IMG_INSTALLED
+            else None,
             "text/plain": lambda s: Text("".join(s)),
         }
         _rich = {
@@ -176,6 +184,10 @@ def __rich_console__(
         self, console: Console, options: ConsoleOptions
     ) -> RenderResult:
         """Rich display of data display outputs."""
+        self.max_size = (
+            int(options.max_width * 2.5),
+            int(options.max_height * 2.5),
+        )  # arbitrarily choose `2.5` for image size
         yield from self.rich_output
 
     @validator("output_type")