diff --git a/changelog.md b/changelog.md index 62fadfe8..89b5d149 100644 --- a/changelog.md +++ b/changelog.md @@ -7,6 +7,7 @@ Features * Allow file target of `$>` redirection to be quoted. * Display of inline plots returned from `.|` operations. * Don't attempt inline plots in the Windows console. +* Save plots, like parquets, with the `.>` operator. Bug Fixes diff --git a/doc/transforms.md b/doc/transforms.md index c856df02..52338c22 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 permitted + * multiple transform steps are not supported * 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 @@ -84,7 +83,8 @@ the `[dataframe]` section of `~/.myclirc`. ## Saving A query result, transformed `DataFrame`, or transformed `Series` can be -written directly to a Parquet file with the `.>` operator. +written directly to a Parquet file with the `.>` operator. An Altair plot can +also be written to a PNG file with the same operator. Save example: @@ -92,22 +92,22 @@ Save example: SELECT * FROM orders .> orders.parquet; ``` -The `.>` operator must be last, requires a `.parquet` destination, and +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. +reports its destination and row count if appropriate. When `post_redirect_command` is set in `~/.myclirc`, the given command runs -after a successful Parquet save. +after a successful file save. `.>` cannot be combined with the `\x` or `\G` special display terminators. ## Combining -Parquet saves may be combined with dataframe transforms. Again, `.>` must -be the last operator. +Parquet saves may be combined with dataframe transforms. Again, the save +operator `.>` must be the last operator. -Combined transform and save example: +Combined transform and save examples: ```sql SELECT * FROM orders .| df.group_by('customer_id').len() .> customer_counts.parquet; @@ -116,3 +116,7 @@ SELECT * FROM orders .| df.group_by('customer_id').len() .> customer_counts.parq ```sql SELECT * FROM orders .| df['order_id'] .> order_ids.parquet; ``` + +```sql +SELECT * FROM orders .| df['total'].plot.hist() .> total_histogram.png; +``` diff --git a/mycli/main_modes/repl.py b/mycli/main_modes/repl.py index a2090205..e09cd690 100644 --- a/mycli/main_modes/repl.py +++ b/mycli/main_modes/repl.py @@ -750,7 +750,7 @@ def _one_iteration( mycli.explorer_formatter.query = text if polars_transform is not None: assert polars_pipeline is not None - if polars_pipeline.parquet_path is None: + if polars_pipeline.output_path is None: polars_result = run_polars_transform( polars_transform, results, @@ -763,22 +763,22 @@ def _one_iteration( polars_result = run_polars_transform( polars_transform, results, - polars_pipeline.parquet_path, + polars_pipeline.output_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_path is None: if polars_pipeline.output_mode == 'explorer': special.set_explorer_output(True) elif polars_pipeline.output_mode == 'expanded': special.set_expanded_output(True) _output_results(mycli, state, iter([polars_result]), start) - if polars_pipeline.parquet_path is not None: + if polars_pipeline.output_path is not None: special.run_post_redirect_hook( mycli.post_redirect_command, - polars_pipeline.parquet_path, + polars_pipeline.output_path, ) else: _output_results(mycli, state, results, start) diff --git a/mycli/packages/polars_transform.py b/mycli/packages/polars_transform.py index 33966b0a..3b0b6a59 100644 --- a/mycli/packages/polars_transform.py +++ b/mycli/packages/polars_transform.py @@ -23,7 +23,7 @@ class PolarsTransformError(RuntimeError): class PolarsPipeline: sql: str expression: str | None - parquet_path: str | None + output_path: str | None output_mode: OutputMode @@ -37,7 +37,7 @@ class PolarsTransform: def parse_polars_transform(command: str) -> PolarsPipeline | None: - """Parse a SQL statement with optional Polars transform and Parquet output.""" + """Parse a SQL statement with optional Polars transform and file output.""" try: tokens = sqlglot.tokenize(command) except sqlglot.errors.TokenError as exc: @@ -56,7 +56,7 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None: if following.end + 1 >= len(command): if following.token_type == sqlglot.TokenType.PIPE: raise PolarsTransformError('Polars transforms require a Python expression.') - raise PolarsTransformError('Parquet saves require a destination path.') + raise PolarsTransformError('File saves require a destination path.') if not command[following.end + 1].isspace(): continue if following.token_type == sqlglot.TokenType.PIPE: @@ -65,7 +65,7 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None: pipe_index = index else: if parquet_index is not None: - raise PolarsTransformError('Parquet saves support only one ".>" operator.') + raise PolarsTransformError('File saves support only one ".>" operator.') parquet_index = index if pipe_index is None and parquet_index is None: @@ -77,20 +77,20 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None: assert first_index is not None sql = command[: tokens[first_index].start].strip() expression: str | None = None - parquet_path: str | None = None + output_path: str | None = None if pipe_index is not None: pipe_operator = tokens[pipe_index + 1] expression_end = tokens[parquet_index].start if parquet_index is not None else len(command) expression = command[pipe_operator.end + 1 : expression_end].strip() if parquet_index is not None: parquet_operator = tokens[parquet_index + 1] - parquet_path = command[parquet_operator.end + 1 :].strip() - parquet_path = parquet_path.removesuffix(delimiter_command.current).rstrip() - parquet_path = parquet_path.removesuffix(r'\g').rstrip() + output_path = command[parquet_operator.end + 1 :].strip() + output_path = output_path.removesuffix(delimiter_command.current).rstrip() + output_path = output_path.removesuffix(r'\g').rstrip() - has_display_terminator = any(value is not None and value.endswith((r'\x', r'\G')) for value in (sql, expression, parquet_path)) - if parquet_path is not None and has_display_terminator: - raise PolarsTransformError('Parquet saves cannot use special display terminators.') + has_display_terminator = any(value is not None and value.endswith((r'\x', r'\G')) for value in (sql, expression, output_path)) + if output_path is not None and has_display_terminator: + raise PolarsTransformError('File saves cannot use special display terminators.') if sql.endswith(r'\x') or expression is not None and expression.endswith(r'\x'): output_mode: OutputMode = 'explorer' elif sql.endswith(r'\G') or expression is not None and expression.endswith(r'\G'): @@ -107,23 +107,23 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None: raise PolarsTransformError('Polars transforms require a SQL statement.') if expression is not None and not expression: raise PolarsTransformError('Polars transforms require a Python expression.') - if parquet_path is not None: - if not parquet_path: - raise PolarsTransformError('Parquet saves require a destination path.') - parquet_path = _parse_parquet_path(parquet_path) + if output_path is not None: + if not output_path: + raise PolarsTransformError('File saves require a destination path.') + output_path = _parse_output_path(output_path) _validate_sql(sql) - return PolarsPipeline(sql=sql, expression=expression, parquet_path=parquet_path, output_mode=output_mode) + return PolarsPipeline(sql=sql, expression=expression, output_path=output_path, output_mode=output_mode) -def _parse_parquet_path(path: str) -> str: +def _parse_output_path(path: str) -> str: if path[0] in ('\'', '"'): if len(path) < 2 or path[-1] != path[0]: - raise PolarsTransformError('Parquet save paths must use matching quotes.') + raise PolarsTransformError('File save paths must use matching quotes.') path = path[1:-1] elif any(character.isspace() for character in path): - raise PolarsTransformError('Parquet save paths containing spaces must be quoted.') - if not path.lower().endswith('.parquet'): - raise PolarsTransformError('Parquet save paths must end in ".parquet".') + 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".') return path @@ -181,7 +181,7 @@ def _load_vl_convert() -> None: def run_polars_transform( transform: PolarsTransform, results: Iterable[SQLResult], - parquet_path: str | None = None, + output_path: str | None = None, *, image_protocol: ImageProtocol = 'none', plot_scale_factor: float = 1.0, @@ -212,34 +212,49 @@ def run_polars_transform( except Exception as exc: raise PolarsTransformError(f'Polars expression failed: {type(exc).__name__}: {exc}') from exc if isinstance(value, transform.polars.DataFrame): - if parquet_path is not None: + if output_path is not None: + if not output_path.lower().endswith('.parquet'): + raise PolarsTransformError('Polars DataFrame results can only be written to ".parquet" files.') try: - value.write_parquet(parquet_path) + value.write_parquet(output_path) except Exception as exc: - raise PolarsTransformError(f'Unable to write Parquet file "{parquet_path}": {type(exc).__name__}: {exc}') from exc - return SQLResult(status=f'Wrote {len(value)} rows to {parquet_path}.') + raise PolarsTransformError(f'Unable to write Parquet file "{output_path}": {type(exc).__name__}: {exc}') from exc + return SQLResult(status=f'Wrote {len(value)} rows to {output_path}.') return SQLResult(header=list(value.columns), rows=list(value.iter_rows())) if isinstance(value, transform.polars.Series): column_name = value.name or 'value' - if parquet_path is not None: + if output_path is not None: + if not output_path.lower().endswith('.parquet'): + raise PolarsTransformError('Polars Series results can only be written to ".parquet" files.') try: series_dataframe = value.rename(column_name).to_frame() - series_dataframe.write_parquet(parquet_path) + series_dataframe.write_parquet(output_path) except Exception as exc: - 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}.') + raise PolarsTransformError(f'Unable to write Parquet file "{output_path}": {type(exc).__name__}: {exc}') from exc + return SQLResult(status=f'Wrote {len(series_dataframe)} rows to {output_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': + if output_path is not None and not output_path.lower().endswith('.png'): + raise PolarsTransformError('Altair charts can only be written to ".png" 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() - 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 + if output_path is not None: + try: + value.save( + output_path, + format='png', + scale_factor=plot_scale_factor, + ppi=plot_ppi, + ) + 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}.') + png = BytesIO() try: value.save( png, @@ -250,6 +265,8 @@ def run_polars_transform( 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.') + if output_path is not None: + raise PolarsTransformError( + 'Polars transforms must return a DataFrame or Series for Parquet output, or an Altair chart for PNG output.' + ) return SQLResult(status=f'Nothing could be displayed for return type: {type(value)}') diff --git a/test/pytests/test_main_modes_repl.py b/test/pytests/test_main_modes_repl.py index d17eadeb..66d0c183 100644 --- a/test/pytests/test_main_modes_repl.py +++ b/test/pytests/test_main_modes_repl.py @@ -1374,6 +1374,61 @@ def run( assert hook_calls == [('post {}', 'orders.parquet')] +def test_one_iteration_writes_polars_png_and_runs_post_redirect_hook(monkeypatch: pytest.MonkeyPatch) -> None: + patch_repl_runtime_defaults(monkeypatch) + + class FakeSQLExecute: + dbname = 'db' + connection_id = 0 + + def run(self, text: str) -> Iterator[SQLResult]: + assert text == 'SELECT * FROM orders' + return iter([SQLResult(header=['id'], rows=[(1,)])]) + + cli = make_repl_cli(FakeSQLExecute()) + cli.post_redirect_command = 'post {}' + transform = object() + run_calls: list[str] = [] + hook_calls: list[tuple[str, str]] = [] + + monkeypatch.setattr(repl_mode, 'prepare_polars_transform', lambda sql, expression: transform) + + def run( + received_transform: object, + results: Iterator[SQLResult], + path: str, + *, + image_protocol: str, + plot_scale_factor: float, + plot_ppi: int, + plot_theme: str, + ) -> SQLResult: + assert received_transform is transform + assert list(results) == [SQLResult(header=['id'], rows=[(1,)])] + assert image_protocol == 'none' + assert plot_scale_factor == 1.0 + assert plot_ppi == 200 + assert plot_theme == 'carbong90' + run_calls.append(path) + return SQLResult(status=f'Wrote PNG image to {path}.') + + monkeypatch.setattr(repl_mode, 'run_polars_transform', run) + monkeypatch.setattr( + repl_mode.special, + 'run_post_redirect_hook', + lambda command, filename: hook_calls.append((command, filename)), + ) + + repl_mode._one_iteration( + cli, + repl_mode.ReplState(), + 'SELECT * FROM orders .| alt.Chart(df) .> orders.png', + ) + + assert run_calls == ['orders.png'] + assert hook_calls == [('post {}', 'orders.png')] + + def test_one_iteration_reports_polars_post_redirect_hook_error(monkeypatch: pytest.MonkeyPatch) -> None: patch_repl_runtime_defaults(monkeypatch) diff --git a/test/pytests/test_polars_transform.py b/test/pytests/test_polars_transform.py index 0ff4ee8b..89b0b351 100644 --- a/test/pytests/test_polars_transform.py +++ b/test/pytests/test_polars_transform.py @@ -75,6 +75,7 @@ class FakeTopLevelMixin: class FakeChart(FakeTopLevelMixin): scale_factors: list[float] = [] ppis: list[int] = [] + saved_paths: list[str] = [] def __init__(self, _data: FakeDataFrame) -> None: self.save_calls: list[str] = [] @@ -83,7 +84,10 @@ 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']) - file.write(b'png image') + if isinstance(file, str): + self.saved_paths.append(file) + else: + file.write(b'png image') class FailingChart(FakeTopLevelMixin): @@ -131,7 +135,7 @@ def test_parse_polars_transform_splits_sql_and_preserves_expression() -> None: assert parse_polars_transform("SELECT * FROM orders .| df.filter(pl.col('state') == 'open')") == PolarsPipeline( sql='SELECT * FROM orders', expression="df.filter(pl.col('state') == 'open')", - parquet_path=None, + output_path=None, output_mode='tabular', ) @@ -148,7 +152,7 @@ def test_parse_polars_transform_returns_output_mode(command: str, output_mode: O assert parse_polars_transform(command) == PolarsPipeline( sql='SELECT * FROM orders', expression='df', - parquet_path=None, + output_path=None, output_mode=output_mode, ) @@ -161,9 +165,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'), ], ) -def test_parse_polars_transform_parses_parquet_redirect( +def test_parse_polars_transform_parses_file_redirect( command: str, expression: str | None, path: str, @@ -171,7 +177,7 @@ def test_parse_polars_transform_parses_parquet_redirect( assert parse_polars_transform(command) == PolarsPipeline( sql='SELECT * FROM orders', expression=expression, - parquet_path=path, + output_path=path, output_mode='tabular', ) @@ -199,7 +205,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"'), + ('SELECT 1 .> export.csv', 'end in ".parquet" or ".png"'), ('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'), @@ -234,9 +240,9 @@ def fail_parse(sql: str, *, read: str) -> list[Any]: parse_polars_transform('SELECT 1 .| df') -def test_parse_parquet_path_rejects_mismatched_quotes() -> None: +def test_parse_output_path_rejects_mismatched_quotes() -> None: with pytest.raises(PolarsTransformError, match='matching quotes'): - polars_transform._parse_parquet_path("'orders.parquet") + polars_transform._parse_output_path("'orders.parquet") def test_prepare_polars_transform_compiles_expression_and_loads_polars(monkeypatch: pytest.MonkeyPatch) -> None: @@ -475,6 +481,50 @@ def test_run_polars_transform_renders_altair_chart_for_image_protocol( 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: + load_calls: list[None] = [] + FakeChart.scale_factors = [] + FakeChart.ppis = [] + FakeChart.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)'), + iter([SQLResult(header=['id'], rows=[(1,)])]), + 'chart.png', + 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 FakeTheme.enabled == ['dark'] + + +def test_run_polars_transform_reports_altair_png_write_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'), + polars=FakePolars, + altair=FailingAltair, + ) + + with pytest.raises(PolarsTransformError, match='Unable to write PNG file "chart.png": OSError: renderer failed'): + run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,)])]), + 'chart.png', + ) + + def test_run_polars_transform_uses_integer_default_plot_ppi(monkeypatch: pytest.MonkeyPatch) -> None: FakeChart.ppis = [] monkeypatch.setattr(polars_transform, '_load_vl_convert', lambda: None) @@ -538,7 +588,7 @@ def fail_load() -> None: def test_run_polars_transform_rejects_altair_chart_parquet_output() -> None: - with pytest.raises(PolarsTransformError, match='must return a DataFrame or Series'): + with pytest.raises(PolarsTransformError, match='Altair charts can only be written'): run_polars_transform( make_transform('alt.Chart(df)'), iter([SQLResult(header=['id'], rows=[(1,)])]), @@ -547,6 +597,22 @@ def test_run_polars_transform_rejects_altair_chart_parquet_output() -> None: ) +@pytest.mark.parametrize( + ('expression', 'message'), + [ + ('df', 'DataFrame results can only be written'), + ("pl.Series('id', [1])", 'Series results can only be written'), + ], +) +def test_run_polars_transform_rejects_dataframe_and_series_png_output(expression: str, message: str) -> None: + with pytest.raises(PolarsTransformError, match=message): + run_polars_transform( + make_transform(expression), + iter([SQLResult(header=['id'], rows=[(1,)])]), + 'data.png', + ) + + def test_prepare_polars_transform_supports_direct_parquet_output(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr('mycli.packages.polars_transform._load_polars', lambda: FakePolars)