From 3fd60b0282ddad7a01530ee71fbcf410da1c5d15 Mon Sep 17 00:00:00 2001 From: Roland Walker Date: Fri, 24 Jul 2026 16:08:34 -0400 Subject: [PATCH] infer plot image-save format from file extension supporting PNG, PDF, SVG, and HTML formats. Incidentally make the choice of "plot" over "chart" more consistent in variable names, messages, etc. --- doc/transforms.md | 17 +-- mycli/packages/polars_transform.py | 42 +++++--- test/pytests/test_main_modes_repl.py | 14 ++- test/pytests/test_polars_transform.py | 145 ++++++++++++++++---------- 4 files changed, 139 insertions(+), 79 deletions(-) diff --git a/doc/transforms.md b/doc/transforms.md index bf0bc2d5..3a979940 100644 --- a/doc/transforms.md +++ b/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,7 +86,7 @@ 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: @@ -95,10 +94,14 @@ Save example: 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. diff --git a/mycli/packages/polars_transform.py b/mycli/packages/polars_transform.py index 0c8f51ba..b623c8d0 100644 --- a/mycli/packages/polars_transform.py +++ b/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() diff --git a/test/pytests/test_main_modes_repl.py b/test/pytests/test_main_modes_repl.py index 66d0c183..78b728fc 100644 --- a/test/pytests/test_main_modes_repl.py +++ b/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: diff --git a/test/pytests/test_polars_transform.py b/test/pytests/test_polars_transform.py index 65f1f085..a25cd2ab 100644 --- a/test/pytests/test_polars_transform.py +++ b/test/pytests/test_polars_transform.py @@ -127,7 +127,8 @@ class FakeTopLevelMixin: pass -class FakeChart(FakeTopLevelMixin): +class FakePlot(FakeTopLevelMixin): + formats: list[str] = [] scale_factors: list[float] = [] ppis: list[int] = [] saved_paths: list[str] = [] @@ -137,15 +138,18 @@ def __init__(self, _data: FakeDataFrame) -> None: def save(self, file: Any, **kwargs: Any) -> None: self.save_calls.append(kwargs['format']) - self.scale_factors.append(float(kwargs['scale_factor'])) - self.ppis.append(kwargs['ppi']) + self.formats.append(kwargs['format']) + if 'scale_factor' in kwargs: + self.scale_factors.append(float(kwargs['scale_factor'])) + if 'ppi' in kwargs: + self.ppis.append(kwargs['ppi']) if isinstance(file, str): self.saved_paths.append(file) else: file.write(b'png image') -class FailingChart(FakeTopLevelMixin): +class FailingPlot(FakeTopLevelMixin): def save(self, _file: Any, **_kwargs: str) -> None: raise OSError('renderer failed') @@ -163,8 +167,8 @@ class FakeAltair: theme = FakeTheme @staticmethod - def Chart(data: FakeDataFrame) -> FakeChart: - return FakeChart(data) + def Plot(data: FakeDataFrame) -> FakePlot: + return FakePlot(data) class FailingAltair: @@ -172,8 +176,8 @@ class FailingAltair: theme = FakeTheme @staticmethod - def Chart(_data: FakeDataFrame) -> FailingChart: - return FailingChart() + def Plot(_data: FakeDataFrame) -> FailingPlot: + return FailingPlot() def make_transform(expression: str) -> PolarsTransform: @@ -220,8 +224,11 @@ def test_parse_polars_transform_returns_output_mode(command: str, output_mode: O (r'SELECT * FROM orders .> orders.parquet \g', None, 'orders.parquet'), ("SELECT * FROM orders .> 'order exports.parquet'", None, 'order exports.parquet'), ('SELECT * FROM orders .| df.head(10) .> orders.parquet', 'df.head(10)', 'orders.parquet'), - ('SELECT * FROM orders .| alt.Chart(df) .> orders.png', 'alt.Chart(df)', 'orders.png'), - (r"SELECT * FROM orders .| alt.Chart(df) .> 'order charts.png' \g", 'alt.Chart(df)', 'order charts.png'), + ('SELECT * FROM orders .| alt.Plot(df) .> orders.png', 'alt.Plot(df)', 'orders.png'), + (r"SELECT * FROM orders .| alt.Plot(df) .> 'order plot.png' \g", 'alt.Plot(df)', 'order plot.png'), + ('SELECT * FROM orders .| alt.Plot(df) .> orders.pdf', 'alt.Plot(df)', 'orders.pdf'), + (r"SELECT * FROM orders .| alt.Plot(df) .> 'order plot.SVG' \g", 'alt.Plot(df)', 'order plot.SVG'), + (r"SELECT * FROM orders .| alt.Plot(df) .> 'order plot.HTML' \g", 'alt.Plot(df)', 'order plot.HTML'), ], ) def test_parse_polars_transform_parses_file_redirect( @@ -260,7 +267,7 @@ def test_parse_polars_transform_ignores_non_suffix_markers(command: str) -> None ('SELECT 1 .|', 'require a Python expression'), ('SELECT 1; SELECT 2 .| df', 'exactly one SQL statement'), ('SELECT 1 .>', 'require a destination path'), - ('SELECT 1 .> export.csv', 'end in ".parquet" or ".png"'), + ('SELECT 1 .> export.csv', 'end in ".parquet", ".png", ".pdf", ".svg", or ".html"'), ('SELECT 1 .> export file.parquet', 'must be quoted'), ('SELECT 1 .> export.parquet .| df', 'must follow'), ('SELECT 1 .| df .> export.parquet \\x', 'cannot use special display terminators'), @@ -634,9 +641,9 @@ def test_run_polars_transform_reports_parquet_write_error() -> None: ) -def test_run_polars_transform_reports_disabled_altair_chart_output() -> None: +def test_run_polars_transform_reports_disabled_altair_plot_output() -> None: result = run_polars_transform( - make_transform('alt.Chart(df)'), + make_transform('alt.Plot(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), ) @@ -644,18 +651,18 @@ def test_run_polars_transform_reports_disabled_altair_chart_output() -> None: @pytest.mark.parametrize('image_protocol', ['iterm2', 'kitty']) -def test_run_polars_transform_renders_altair_chart_for_image_protocol( +def test_run_polars_transform_renders_altair_plot_for_image_protocol( monkeypatch: pytest.MonkeyPatch, image_protocol: ImageProtocol, ) -> None: load_calls: list[None] = [] - FakeChart.scale_factors = [] - FakeChart.ppis = [] + FakePlot.scale_factors = [] + FakePlot.ppis = [] FakeTheme.enabled = [] monkeypatch.setattr(polars_transform, '_load_vl_convert', lambda: load_calls.append(None)) result = run_polars_transform( - make_transform('alt.Chart(df)'), + make_transform('alt.Plot(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), image_protocol=image_protocol, plot_scale_factor=1.5, @@ -664,67 +671,94 @@ def test_run_polars_transform_renders_altair_chart_for_image_protocol( ) assert load_calls == [None] - assert FakeChart.scale_factors == [1.5] - assert FakeChart.ppis == [144] + assert FakePlot.scale_factors == [1.5] + assert FakePlot.ppis == [144] assert FakeTheme.enabled == ['dark'] assert result == SQLResult(image=b'png image', image_protocol=image_protocol) -def test_run_polars_transform_writes_altair_chart_to_png(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize( + ('path', 'plot_format', 'expected_scale_factors', 'expected_ppis', 'expected_load_calls', 'output_kind'), + [ + ('plot.png', 'png', [1.5], [144], [None], 'image'), + ('plot.pdf', 'pdf', [1.5], [], [None], 'image'), + ('plot.SVG', 'svg', [1.5], [], [None], 'image'), + ('plot.html', 'html', [], [], [], 'document'), + ], +) +def test_run_polars_transform_writes_altair_plot_to_file( + monkeypatch: pytest.MonkeyPatch, + path: str, + plot_format: str, + expected_scale_factors: list[float], + expected_ppis: list[int], + expected_load_calls: list[None], + output_kind: str, +) -> None: load_calls: list[None] = [] - FakeChart.scale_factors = [] - FakeChart.ppis = [] - FakeChart.saved_paths = [] + FakePlot.formats = [] + FakePlot.scale_factors = [] + FakePlot.ppis = [] + FakePlot.saved_paths = [] FakeTheme.enabled = [] monkeypatch.setattr(polars_transform, '_load_vl_convert', lambda: load_calls.append(None)) result = run_polars_transform( - make_transform('alt.Chart(df)'), + make_transform('alt.Plot(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), - 'chart.png', + path, image_protocol='none', plot_scale_factor=1.5, plot_ppi=144, plot_theme='dark', ) - assert result == SQLResult(status='Wrote PNG image to chart.png.') - assert load_calls == [None] - assert FakeChart.saved_paths == ['chart.png'] - assert FakeChart.scale_factors == [1.5] - assert FakeChart.ppis == [144] + assert result == SQLResult(status=f'Wrote {plot_format.upper()} {output_kind} to {path}.') + assert load_calls == expected_load_calls + assert FakePlot.formats == [plot_format] + assert FakePlot.saved_paths == [path] + assert FakePlot.scale_factors == expected_scale_factors + assert FakePlot.ppis == expected_ppis assert FakeTheme.enabled == ['dark'] -def test_run_polars_transform_reports_altair_png_write_error(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize('plot_format', ['png', 'pdf', 'svg', 'html']) +def test_run_polars_transform_reports_altair_file_write_error( + monkeypatch: pytest.MonkeyPatch, + plot_format: str, +) -> None: monkeypatch.setattr(polars_transform, '_load_vl_convert', lambda: None) transform = PolarsTransform( sql='SELECT id FROM orders', - expression='alt.Chart(df)', - code=compile('alt.Chart(df)', '', 'eval'), + expression='alt.Plot(df)', + code=compile('alt.Plot(df)', '', 'eval'), polars=FakePolars, altair=FailingAltair, ) - with pytest.raises(PolarsTransformError, match='Unable to write PNG file "chart.png": OSError: renderer failed'): + path = f'plot.{plot_format}' + with pytest.raises( + PolarsTransformError, + match=f'Unable to write {plot_format.upper()} file "{path}": OSError: renderer failed', + ): run_polars_transform( transform, iter([SQLResult(header=['id'], rows=[(1,)])]), - 'chart.png', + path, ) def test_run_polars_transform_uses_integer_default_plot_ppi(monkeypatch: pytest.MonkeyPatch) -> None: - FakeChart.ppis = [] + FakePlot.ppis = [] monkeypatch.setattr(polars_transform, '_load_vl_convert', lambda: None) run_polars_transform( - make_transform('alt.Chart(df)'), + make_transform('alt.Plot(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), image_protocol='kitty', ) - assert FakeChart.ppis == [200] + assert FakePlot.ppis == [200] def test_run_polars_transform_reports_invalid_altair_theme(monkeypatch: pytest.MonkeyPatch) -> None: @@ -737,24 +771,24 @@ def fail_enable(_name: str) -> None: with pytest.raises(PolarsTransformError, match='Unable to enable Altair plot theme "not-a-theme": ValueError: unknown theme'): run_polars_transform( - make_transform('alt.Chart(df)'), + make_transform('alt.Plot(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), image_protocol='kitty', plot_theme='not-a-theme', ) -def test_run_polars_transform_reports_altair_chart_render_error(monkeypatch: pytest.MonkeyPatch) -> None: +def test_run_polars_transform_reports_altair_plot_render_error(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(polars_transform, '_load_vl_convert', lambda: None) transform = PolarsTransform( sql='SELECT id FROM orders', - expression='alt.Chart(df)', - code=compile('alt.Chart(df)', '', 'eval'), + expression='alt.Plot(df)', + code=compile('alt.Plot(df)', '', 'eval'), polars=FakePolars, altair=FailingAltair, ) - with pytest.raises(PolarsTransformError, match='Unable to render Altair chart: OSError: renderer failed'): + with pytest.raises(PolarsTransformError, match='Unable to render Altair plot: OSError: renderer failed'): run_polars_transform( transform, iter([SQLResult(header=['id'], rows=[(1,)])]), @@ -770,35 +804,40 @@ def fail_load() -> None: with pytest.raises(PolarsTransformError, match='requires vl-convert-python'): run_polars_transform( - make_transform('alt.Chart(df)'), + make_transform('alt.Plot(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), image_protocol='iterm2', ) -def test_run_polars_transform_rejects_altair_chart_parquet_output() -> None: - with pytest.raises(PolarsTransformError, match='Altair charts can only be written'): +def test_run_polars_transform_rejects_altair_plot_parquet_output() -> None: + with pytest.raises(PolarsTransformError, match='Altair plots can only be written'): run_polars_transform( - make_transform('alt.Chart(df)'), + make_transform('alt.Plot(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), - 'chart.parquet', + 'plot.parquet', image_protocol='iterm2', ) @pytest.mark.parametrize( - ('expression', 'message'), + ('expression', 'path', 'message'), [ - ('df', 'DataFrame results can only be written'), - ("pl.Series('id', [1])", 'Series results can only be written'), + ('df', 'data.pdf', 'DataFrame results can only be written'), + ("pl.Series('id', [1])", 'data.svg', 'Series results can only be written'), + ('df', 'data.html', 'DataFrame results can only be written'), ], ) -def test_run_polars_transform_rejects_dataframe_and_series_png_output(expression: str, message: str) -> None: +def test_run_polars_transform_rejects_dataframe_and_series_plot_output( + expression: str, + path: str, + message: str, +) -> None: with pytest.raises(PolarsTransformError, match=message): run_polars_transform( make_transform(expression), iter([SQLResult(header=['id'], rows=[(1,)])]), - 'data.png', + path, )