ability to display static Altair plots

* if a .| dataframe operation returns a plot, display it inline, using
  the kitty graphics protocol by default (since it has the widest
  support in contemporary terminals such as kitty and Ghostty
  https://sw.kovidgoyal.net/kitty/graphics-protocol/)
* revise transforms.md doc, including changes outside the subject area
  of plotting

There are some small incidental PromptState tests added, as "coverage"
showed that our coverage there was not 100% for some reason.

There are many minor followups we can imagine here, such as providing
the facility to save an image.

Author: rolandwalker

README.md

@@ -71,7 +71,7 @@ Features
 * Pretty print tabular data (with colors!).
 * Support for SSL connections
 * Shell-style trailing redirects with `$>`, `$>>` and `$|` operators.
-* [Polars](https://pola.rs) dataframe [transforms](doc/transforms.md) with `.|` and Parquet saves with `.>`.
+* [Polars](https://pola.rs) dataframe [transforms and plots](doc/transforms.md) with `.|`, and Parquet saves with `.>`.
 * Support for querying LLMs with context derived from your schema using `/llm`.
 * Support for storing passwords in the system keyring.
 

changelog.md

@@ -5,6 +5,7 @@ Features
 ---------
 * Subcommand completions for the `/dsn` command.
 * Allow file target of `$>` redirection to be quoted.
+* Display of inline plots returned from `.|` operations.
 
 
 Bug Fixes

doc/screenshots/total_histogram.png

@@ -1,4 +1,4 @@
-# Transforms with Polars Dataframes
+# Transforms, Plots, and Parquets with Polars Dataframes
 
 ## Installing
 
@@ -8,19 +8,24 @@ Install mycli with dataframe support using:
 pip install --upgrade 'mycli[dataframe]'
 ```
 
-or install the Polars and Altair libraries separately.
+or install these libraries separately:
+
+ * [`polars`](https://pypi.org/project/polars/)
+ * [`altair`](https://pypi.org/project/altair/)
+ * [`vl-convert-python`](https://pypi.org/project/vl-convert-python/)
 
 ## BETA STATUS
 
-Dataframe transforms are new and experimental.  The interface and functionality
-may still change.
+Dataframe transforms and plots are new and experimental.  The interface and
+functionality may still change.
 
 Here are some known limitations:
 
- * transforms can't be mixed with `$|` shell redirection
+ * transforms can't be composed with `$|` shell redirection
  * multiple transform steps are not permitted
- * the Altair library is provided, but plots can't yet be displayed
  * results from `UNION`s may be unable to be transformed
+ * images cannot yet be saved
+ * PNG images are static and do not support all Altair features
 
 And there are inherent limitations to the post-processing model: the entire
 SQL result must be transferred from the server and loaded into local memory.
@@ -34,7 +39,7 @@ single SQL statement.  The Python expression receives
  * `pl`, the Polars module
  * `alt`, the Altair module
 
-Spaces may be required around the operator.
+Spaces may be required around the `.|` operator.
 
 Transform example:
 
@@ -51,20 +56,34 @@ SELECT customer_id, COUNT(1) AS len FROM orders GROUP BY customer_id;
 Transform expressions run with normal Python privileges, and expressions
 should not be run from untrusted sources.  If the transform operation
 returns a Polars `DataFrame` or `Series`, the result is rendered by mycli
-as tabular output; other return types currently give a warning and cannot
-be displayed.
+as tabular output.  Most other return types will be silently ignored.
 
 Transform expressions are useful for operations such as medians which
-cannot be done (or are simply awkward) in SQL.  Example:
+cannot be done (or are awkward) in SQL.  Example:
 
 ```sql
 SELECT * FROM orders .| df.describe();
 ```
 
+## Plotting
+
+If the dataframe transform operation returns an Altair plot, the result can
+be rendered as an inline PNG in many terminals.
+
+Example:
+
+```sql
+SELECT * FROM orders .| df['total'].plot.hist();
+```
+![histogram](https://raw.githubusercontent.com/dbcli/mycli/main/doc/screenshots/total_histogram.png)
+
+Image size, display protocol, and other properties can be configured in
+the `[dataframe]` section of `~/.myclirc`.
+
 ## Saving
 
 A query result, transformed `DataFrame`, or transformed `Series` can be
-written directly to Parquet with the `.>` operator.
+written directly to a Parquet file with the `.>` operator.
 
 Save example:
 

doc/transforms.md

@@ -6,6 +6,7 @@
 from configobj import ConfigObj
 
 from mycli.config import strip_matching_quotes
+from mycli.types import ImageProtocol
 
 if TYPE_CHECKING:
     from mycli.client import MyCli
@@ -46,6 +47,16 @@ def normalize_ssl_mode(
     return ssl_mode, error_notice
 
 
+def normalize_image_protocol(image_protocol: str | None) -> tuple[ImageProtocol, str | None]:
+    if image_protocol == 'iterm2':
+        return 'iterm2', None
+    if image_protocol == 'kitty':
+        return 'kitty', None
+    if image_protocol in ('none', '', None):
+        return 'none', None
+    return 'none', f'Invalid config option provided for image_protocol ({image_protocol}); disabling.'
+
+
 def configure_prompt_state(
     mycli: MyCli,
     config: ConfigObj,

mycli/app_state.py

@@ -17,6 +17,7 @@
     configure_prompt_state,
     destructive_keywords_from_config,
     llm_prompt_truncation,
+    normalize_image_protocol,
     normalize_ssl_mode,
 )
 from mycli.client_commands import ClientCommandsMixin
@@ -145,6 +146,12 @@ def __init__(
         self.null_string = c['main'].get('null_string')
         self.numeric_alignment = c['main'].get('numeric_alignment', 'right') or 'right'
         self.binary_display = c['main'].get('binary_display')
+        self.image_protocol, image_protocol_error = normalize_image_protocol(c['dataframe'].get('image_protocol'))
+        if image_protocol_error:
+            self.echo(image_protocol_error, err=True, fg='red')
+        self.plot_scale_factor = c['dataframe'].as_float('plot_scale_factor')
+        self.plot_ppi = c['dataframe'].as_int('plot_ppi')
+        self.plot_theme = c['dataframe'].get('plot_theme', 'carbong90') or 'carbong90'
         self.llm_prompt_field_truncate, self.llm_prompt_section_truncate = llm_prompt_truncation(c)
 
         self.ssl_mode, ssl_mode_error = normalize_ssl_mode(c, self.config_without_package_defaults)

mycli/client.py

@@ -751,9 +751,24 @@ def _one_iteration(
         if polars_transform is not None:
             assert polars_pipeline is not None
             if polars_pipeline.parquet_path is None:
-                polars_result = run_polars_transform(polars_transform, results)
+                polars_result = run_polars_transform(
+                    polars_transform,
+                    results,
+                    image_protocol=mycli.image_protocol,
+                    plot_scale_factor=mycli.plot_scale_factor,
+                    plot_ppi=mycli.plot_ppi,
+                    plot_theme=mycli.plot_theme,
+                )
             else:
-                polars_result = run_polars_transform(polars_transform, results, polars_pipeline.parquet_path)
+                polars_result = run_polars_transform(
+                    polars_transform,
+                    results,
+                    polars_pipeline.parquet_path,
+                    image_protocol=mycli.image_protocol,
+                    plot_scale_factor=mycli.plot_scale_factor,
+                    plot_ppi=mycli.plot_ppi,
+                    plot_theme=mycli.plot_theme,
+                )
             if polars_pipeline.parquet_path is None:
                 if polars_pipeline.output_mode == 'explorer':
                     special.set_explorer_output(True)

mycli/main_modes/repl.py

@@ -265,6 +265,27 @@ format = mysql_unicode
 # Whether to remove the last line from the formatted output.
 trim_footer = False
 
+[dataframe]
+
+# How to display inline image results in a terminal emulator. Possible values:
+# * kitty
+# * iterm2
+# empty to disable
+image_protocol = kitty
+
+# PNG plot resolution in pixels per inch. Must be a positive integer.
+# The apparent size is a combination of plot_ppi and plot_scale_factor.
+plot_ppi = 200
+
+# PNG plot resolution multiplier. Must be a positive number.
+# The apparent size is a combination of plot_ppi and plot_scale_factor.
+plot_scale_factor = 1.0
+
+# Altair theme for rendering plots. Examples: carbong90, dark, default.
+# Available themes depend on the installed libraries.  See
+# https://github.com/vega/vega-themes/#included-themes
+plot_theme = carbong90
+
 [search]
 
 # Whether to apply syntax highlighting to the preview window in fuzzy history

mycli/myclirc

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import base64
 from datetime import datetime
 from decimal import Decimal
 from io import TextIOWrapper
@@ -115,6 +116,13 @@ def output(
         is_warnings_style: bool = False,
     ) -> None:
         """Output text to stdout or a pager command."""
+        if result.image is not None:
+            if result.image_protocol == 'iterm2':
+                click.secho('')
+                self.output_iterm2_image(result.image)
+            elif result.image_protocol == 'kitty':
+                click.secho('')
+                self.output_kitty_image(result.image)
         if output:
             if self.prompt_session is not None:
                 size = self.prompt_session.output.get_size()
@@ -186,6 +194,22 @@ def newlinewrapper(text: list[str]) -> Generator[str, None, None]:
             styled_status = to_formatted_text(status, style=add_style)
             prompt_toolkit.print_formatted_text(styled_status, style=self.ptoolkit_style)
 
+    def output_iterm2_image(self, image: bytes) -> None:
+        """Emit a PNG using the iTerm2 inline image protocol."""
+        filename = base64.b64encode(b'chart.png').decode('ascii')
+        contents = base64.b64encode(image).decode('ascii')
+        click.echo(f'\x1b]1337;File=name={filename};size={len(image)};width=auto;height=auto;preserveAspectRatio=1;inline=1:{contents}\x07')
+
+    def output_kitty_image(self, image: bytes) -> None:
+        """Emit a PNG using Kitty's direct graphics protocol."""
+        contents = base64.b64encode(image).decode('ascii')
+        chunks = [contents[index : index + 4096] for index in range(0, len(contents), 4096)] or ['']
+        for index, chunk in enumerate(chunks):
+            action = 'a=T,f=100,' if index == 0 else ''
+            more = 1 if index < len(chunks) - 1 else 0
+            click.echo(f'\x1b_G{action}m={more};{chunk}\x1b\\', nl=False)
+        click.echo()
+
     def configure_pager(self) -> None:
         if not os.environ.get("LESS"):
             os.environ["LESS"] = "-RXF"

mycli/output.py

@@ -2,14 +2,15 @@
 
 import builtins
 from dataclasses import dataclass
+from io import BytesIO
 from types import CodeType
 from typing import Any, Iterable
 
 import sqlglot
 
 from mycli.packages.special.delimitercommand import DelimiterCommand
 from mycli.packages.sqlresult import SQLResult
-from mycli.types import OutputMode
+from mycli.types import ImageProtocol, OutputMode
 
 delimiter_command = DelimiterCommand()
 
@@ -169,10 +170,22 @@ def _load_altair() -> Any:
     return alt
 
 
+def _load_vl_convert() -> None:
+    try:
+        import vl_convert  # noqa: F401
+    except ImportError as exc:
+        raise PolarsTransformError('Altair plot rendering requires vl-convert-python. Install mycli[dataframe].') from exc
+
+
 def run_polars_transform(
     transform: PolarsTransform,
     results: Iterable[SQLResult],
     parquet_path: str | None = None,
+    *,
+    image_protocol: ImageProtocol = 'none',
+    plot_scale_factor: float = 1.0,
+    plot_ppi: int = 200,
+    plot_theme: str = 'carbong90',
 ) -> SQLResult:
     iterator = iter(results)
     try:
@@ -215,6 +228,27 @@ def run_polars_transform(
                 raise PolarsTransformError(f'Unable to write Parquet file "{parquet_path}": {type(exc).__name__}: {exc}') from exc
             return SQLResult(status=f'Wrote {len(series_dataframe)} rows to {parquet_path}.')
         return SQLResult(header=[column_name], rows=[(item,) for item in value])
+    if transform.altair is not None and isinstance(value, transform.altair.TopLevelMixin):
+        if parquet_path is not None:
+            raise PolarsTransformError('Polars transforms must return a DataFrame or Series before writing Parquet output.')
+        if image_protocol == 'none':
+            return SQLResult(status='image_protocol is unset in ~/.myclirc. Inline plotting is disabled.')
+        _load_vl_convert()
+        png = BytesIO()
+        try:
+            transform.altair.theme.enable(plot_theme)
+        except Exception as exc:
+            raise PolarsTransformError(f'Unable to enable Altair plot theme "{plot_theme}": {type(exc).__name__}: {exc}') from exc
+        try:
+            value.save(
+                png,
+                format='png',
+                scale_factor=plot_scale_factor,
+                ppi=plot_ppi,
+            )
+        except Exception as exc:
+            raise PolarsTransformError(f'Unable to render Altair chart: {type(exc).__name__}: {exc}') from exc
+        return SQLResult(image=png.getvalue(), image_protocol=image_protocol)
     if parquet_path is not None:
         raise PolarsTransformError('Polars transforms must return a DataFrame or Series before writing Parquet output.')
     return SQLResult(status=f'Nothing could be displayed for return type: {type(value)}')