Infer plot image-save format from file extension of destination

doc/transforms.md

@@ -22,9 +22,8 @@ functionality may still change.
 Here are some known limitations:
 
  * transforms can't be composed with `$|` shell redirection
- * multiple transform steps are not supported
+ * composing multiple transform operations is not supported
  * results from `UNION`s may be unable to be transformed
- * 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.
@@ -87,18 +86,22 @@ the `[dataframe]` section of `~/.myclirc`.
 
 A query result, transformed `DataFrame`, or transformed `Series` can be
 written directly to a Parquet file with the `.>` operator. An Altair plot can
-also be written to a PNG file with the same operator.
+also be written to a file with the same operator.
 
 Save example:
 
 ```sql
 SELECT * FROM orders .> orders.parquet;
 ```
 
-The `.>` operator must be last, requires a `.parquet` or `.png` destination, and
-overwrites any existing file.  Spaces may be required around the operator.
-Destination paths containing whitespace must be quoted.  A successful write
-reports its destination and row count if appropriate.
+The `.>` operator must be last, requires a `.parquet`, `.png`, `.pdf`, `.svg`,
+or `.html` file extension on the destination, and overwrites any existing file.
+
+Spaces may be required around the operator.  Destination paths containing
+whitespace must be quoted.  A successful write reports its destination and row
+count if appropriate.
+
+When saving a plot, the format is deduced from the file extension.
 
 When `post_redirect_command` is set in `~/.myclirc`, the given command runs
 after a successful file save.

mycli/packages/polars_transform.py

@@ -16,6 +16,7 @@
 from mycli.types import ImageProtocol, OutputMode
 
 delimiter_command = DelimiterCommand()
+PLOT_FORMATS = ('png', 'pdf', 'svg', 'html')
 
 
 class PolarsTransformError(RuntimeError):
@@ -144,11 +145,16 @@ def _parse_output_path(path: str) -> str:
         path = path[1:-1]
     elif any(character.isspace() for character in path):
         raise PolarsTransformError('File save paths containing spaces must be quoted.')
-    if not path.lower().endswith(('.parquet', '.png')):
-        raise PolarsTransformError('File save paths must end in ".parquet" or ".png".')
+    if not path.lower().endswith(('.parquet', '.png', '.pdf', '.svg', '.html')):
+        raise PolarsTransformError('File save paths must end in ".parquet", ".png", ".pdf", ".svg", or ".html".')
     return path
 
 
+def _plot_format_for_path(path: str) -> str | None:
+    path = path.lower()
+    return next((plot_format for plot_format in PLOT_FORMATS if path.endswith(f'.{plot_format}')), None)
+
+
 def _validate_sql(sql: str) -> None:
     try:
         statements = sqlglot.parse(sql, read='mysql')
@@ -278,26 +284,34 @@ def run_polars_transform(
             return SQLResult(status=f'Wrote {len(series_dataframe)} rows to {output_path}.')
         return SQLResult(header=[column_name], rows=[(item,) for item in value])
     elif transform.altair is not None and isinstance(value, transform.altair.TopLevelMixin):
-        if output_path is not None and not output_path.lower().endswith('.png'):
-            raise PolarsTransformError('Altair charts can only be written to ".png" files.')
+        plot_format = _plot_format_for_path(output_path) if output_path is not None else None
+        if output_path is not None and plot_format is None:
+            raise PolarsTransformError('Altair plots can only be written to ".png", ".pdf", ".svg", or ".html" files.')
         if output_path is None and image_protocol == 'none':
             return SQLResult(status='image_protocol is unset in ~/.myclirc. Inline plotting is disabled.')
-        _load_vl_convert()
+        if plot_format != 'html':
+            _load_vl_convert()
         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
         if output_path is not None:
+            assert plot_format is not None
+            save_kwargs: dict[str, Any] = {
+                'format': plot_format,
+            }
+            if plot_format != 'html':
+                save_kwargs['scale_factor'] = plot_scale_factor
+            if plot_format == 'png':
+                save_kwargs['ppi'] = plot_ppi
             try:
-                value.save(
-                    output_path,
-                    format='png',
-                    scale_factor=plot_scale_factor,
-                    ppi=plot_ppi,
-                )
+                value.save(output_path, **save_kwargs)
             except Exception as exc:
-                raise PolarsTransformError(f'Unable to write PNG file "{output_path}": {type(exc).__name__}: {exc}') from exc
-            return SQLResult(status=f'Wrote PNG image to {output_path}.')
+                raise PolarsTransformError(
+                    f'Unable to write {plot_format.upper()} file "{output_path}": {type(exc).__name__}: {exc}'
+                ) from exc
+            output_kind = 'document' if plot_format == 'html' else 'image'
+            return SQLResult(status=f'Wrote {plot_format.upper()} {output_kind} to {output_path}.')
         png = BytesIO()
         try:
             value.save(
@@ -307,7 +321,7 @@ def run_polars_transform(
                 ppi=plot_ppi,
             )
         except Exception as exc:
-            raise PolarsTransformError(f'Unable to render Altair chart: {type(exc).__name__}: {exc}') from exc
+            raise PolarsTransformError(f'Unable to render Altair plot: {type(exc).__name__}: {exc}') from exc
         return SQLResult(image=png.getvalue(), image_protocol=image_protocol)
     elif value is None:
         return SQLResult()

test/pytests/test_main_modes_repl.py

@@ -1374,7 +1374,11 @@ def run(
     assert hook_calls == [('post {}', 'orders.parquet')]
 
 
-def test_one_iteration_writes_polars_png_and_runs_post_redirect_hook(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.parametrize('path', ['orders.png', 'orders.pdf', 'orders.svg', 'orders.html'])
+def test_one_iteration_writes_polars_plot_and_runs_post_redirect_hook(
+    monkeypatch: pytest.MonkeyPatch,
+    path: str,
+) -> None:
     patch_repl_runtime_defaults(monkeypatch)
 
     class FakeSQLExecute:
@@ -1410,7 +1414,7 @@ def run(
         assert plot_ppi == 200
         assert plot_theme == 'carbong90'
         run_calls.append(path)
-        return SQLResult(status=f'Wrote PNG image to {path}.')
+        return SQLResult(status=f'Wrote plot to {path}.')
 
     monkeypatch.setattr(repl_mode, 'run_polars_transform', run)
     monkeypatch.setattr(
@@ -1422,11 +1426,11 @@ def run(
     repl_mode._one_iteration(
         cli,
         repl_mode.ReplState(),
-        'SELECT * FROM orders .| alt.Chart(df) .> orders.png',
+        f'SELECT * FROM orders .| alt.Plot(df) .> {path}',
     )
 
-    assert run_calls == ['orders.png']
-    assert hook_calls == [('post {}', 'orders.png')]
+    assert run_calls == [path]
+    assert hook_calls == [('post {}', path)]
 
 
 def test_one_iteration_reports_polars_post_redirect_hook_error(monkeypatch: pytest.MonkeyPatch) -> None: