diff --git a/README.md b/README.md index 9b67059f3..6d3705797 100644 --- a/README.md +++ b/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. diff --git a/changelog.md b/changelog.md index b12ae07ea..5ceb42fb7 100644 --- a/changelog.md +++ b/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 diff --git a/doc/screenshots/total_histogram.png b/doc/screenshots/total_histogram.png new file mode 100644 index 000000000..fe68d1d48 Binary files /dev/null and b/doc/screenshots/total_histogram.png differ diff --git a/doc/transforms.md b/doc/transforms.md index 81779f730..ce8c65775 100644 --- a/doc/transforms.md +++ b/doc/transforms.md @@ -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: diff --git a/mycli/app_state.py b/mycli/app_state.py index d8381ea77..02e175e51 100644 --- a/mycli/app_state.py +++ b/mycli/app_state.py @@ -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, diff --git a/mycli/client.py b/mycli/client.py index 8dd4c3706..e961f15ac 100644 --- a/mycli/client.py +++ b/mycli/client.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) diff --git a/mycli/main_modes/repl.py b/mycli/main_modes/repl.py index 1825c643c..a20902058 100644 --- a/mycli/main_modes/repl.py +++ b/mycli/main_modes/repl.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) diff --git a/mycli/myclirc b/mycli/myclirc index 12ee2c840..a7882745d 100644 --- a/mycli/myclirc +++ b/mycli/myclirc @@ -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 diff --git a/mycli/output.py b/mycli/output.py index 6b3a23204..55b93bb4e 100644 --- a/mycli/output.py +++ b/mycli/output.py @@ -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" diff --git a/mycli/packages/polars_transform.py b/mycli/packages/polars_transform.py index 7a015e243..2c7bc2c7b 100644 --- a/mycli/packages/polars_transform.py +++ b/mycli/packages/polars_transform.py @@ -2,6 +2,7 @@ import builtins from dataclasses import dataclass +from io import BytesIO from types import CodeType from typing import Any, Iterable @@ -9,7 +10,7 @@ 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)}') diff --git a/mycli/packages/sqlresult.py b/mycli/packages/sqlresult.py index b1f5e272b..0b9d4b9c7 100644 --- a/mycli/packages/sqlresult.py +++ b/mycli/packages/sqlresult.py @@ -4,6 +4,8 @@ from prompt_toolkit.formatted_text import FormattedText, to_plain_text from pymysql.cursors import Cursor +from mycli.types import ImageProtocol + @dataclass class SQLResult: @@ -13,9 +15,14 @@ class SQLResult: postamble: str | None = None status: str | FormattedText | None = None command: dict[str, str | float] | None = None + image: bytes | None = None + image_protocol: ImageProtocol = 'none' def __str__(self): - return f"{self.preamble}, {self.header}, {self.rows}, {self.postamble}, {self.status}, {self.command}" + image = f'<{len(self.image)} bytes>' if self.image is not None else None + return ( + f"{self.preamble}, {self.header}, {self.rows}, {self.postamble}, {self.status}, {self.command}, {image}, {self.image_protocol}" + ) @cached_property def status_plain(self): diff --git a/mycli/types.py b/mycli/types.py index 85f2d6d3a..2a00632c6 100644 --- a/mycli/types.py +++ b/mycli/types.py @@ -9,3 +9,9 @@ 'expanded', 'tabular', ] + +ImageProtocol = Literal[ + 'none', + 'iterm2', + 'kitty', +] diff --git a/pyproject.toml b/pyproject.toml index adc9e51f9..3d1da695c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ llm = [ dataframe = [ "polars ~= 1.42.1", "altair ~= 6.2.2", + "vl-convert-python ~= 1.9.0", ] all = [ "mycli[llm,dataframe]", @@ -73,6 +74,7 @@ dev = [ "ruff ~= 0.15.0", "polars ~= 1.42.1", "altair ~= 6.2.2", + "vl-convert-python ~= 1.9.0", ] [project.scripts] diff --git a/test/myclirc b/test/myclirc index e13fa5c7a..fca7a1a5b 100644 --- a/test/myclirc +++ b/test/myclirc @@ -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 diff --git a/test/pytests/test_app_state.py b/test/pytests/test_app_state.py index f97bf9578..94bedcaa7 100644 --- a/test/pytests/test_app_state.py +++ b/test/pytests/test_app_state.py @@ -5,8 +5,10 @@ from mycli.app_state import ( AppStateMixin, + configure_prompt_state, destructive_keywords_from_config, llm_prompt_truncation, + normalize_image_protocol, normalize_ssl_mode, ) @@ -16,6 +18,76 @@ def __init__(self, login_path: str | None = None) -> None: self.login_path = login_path +class PromptState: + prompt_format: str + prompt_lines: int + multiline_continuation_char: str + toolbar_format: str + terminal_tab_title_format: str + terminal_window_title_format: str + multiplex_window_title_format: str + multiplex_pane_title_format: str + + def __init__(self) -> None: + self.default_prompt = 'default> ' + + +@pytest.mark.parametrize( + ('image_protocol', 'expected'), + [ + ('iterm2', ('iterm2', None)), + ('kitty', ('kitty', None)), + ('none', ('none', None)), + ('', ('none', None)), + (None, ('none', None)), + ('unknown', ('none', 'Invalid config option provided for image_protocol (unknown); disabling.')), + ], +) +def test_normalize_image_protocol(image_protocol: str | None, expected: tuple[str, str | None]) -> None: + assert normalize_image_protocol(image_protocol) == expected + + +@pytest.mark.parametrize( + ('prompt', 'config_prompt', 'toolbar_format', 'config_toolbar', 'expected_prompt', 'expected_toolbar'), + [ + ('custom> ', 'configured> ', 'custom toolbar', 'configured toolbar', 'custom> ', 'custom toolbar'), + (None, 'configured> ', None, 'configured toolbar', 'configured> ', 'configured toolbar'), + ('', '', '', 'configured toolbar', 'default> ', 'configured toolbar'), + ], +) +def test_configure_prompt_state_uses_overrides_and_fallbacks( + prompt: str | None, + config_prompt: str, + toolbar_format: str | None, + config_toolbar: str, + expected_prompt: str, + expected_toolbar: str, +) -> None: + state = PromptState() + config = ConfigObj({ + 'main': { + 'prompt': config_prompt, + 'prompt_continuation': '... ', + 'toolbar': config_toolbar, + 'terminal_tab_title': 'tab title', + 'terminal_window_title': 'window title', + 'multiplex_window_title': 'multiplex window title', + 'multiplex_pane_title': 'multiplex pane title', + }, + }) + + configure_prompt_state(state, config, prompt, toolbar_format) # type: ignore[arg-type] + + assert state.prompt_format == expected_prompt + assert state.prompt_lines == 0 + assert state.multiline_continuation_char == '... ' + assert state.toolbar_format == expected_toolbar + assert state.terminal_tab_title_format == 'tab title' + assert state.terminal_window_title_format == 'window title' + assert state.multiplex_window_title_format == 'multiplex window title' + assert state.multiplex_pane_title_format == 'multiplex pane title' + + @pytest.mark.parametrize('ssl_mode', ['auto', 'on', 'off']) def test_normalize_ssl_mode_accepts_known_values(ssl_mode: str) -> None: config = ConfigObj({'main': {'ssl_mode': ssl_mode}, 'connection': {'default_ssl_mode': ssl_mode}}) diff --git a/test/pytests/test_client.py b/test/pytests/test_client.py index 1c5489941..733b112df 100644 --- a/test/pytests/test_client.py +++ b/test/pytests/test_client.py @@ -42,6 +42,76 @@ def test_init_reports_invalid_ssl_mode(monkeypatch: pytest.MonkeyPatch, tmp_path assert echo_calls == [('Invalid config option provided for ssl_mode (invalid); ignoring.', {'err': True, 'fg': 'red'})] +def test_init_configures_image_protocol(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + patch_constructor_side_effects(monkeypatch) + myclirc = write_myclirc( + tmp_path, + """ + [dataframe] + image_protocol = iterm2 + plot_scale_factor = 1.5 + plot_ppi = 144 + plot_theme = dark + """, + ) + + cli = MyCli(myclirc=myclirc) + + assert cli.image_protocol == 'iterm2' + assert cli.plot_scale_factor == 1.5 + assert cli.plot_ppi == 144 + assert isinstance(cli.plot_ppi, int) + assert cli.plot_theme == 'dark' + + +def test_init_uses_default_plot_theme_for_empty_value(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + patch_constructor_side_effects(monkeypatch) + myclirc = write_myclirc( + tmp_path, + """ + [dataframe] + plot_theme = + """, + ) + + cli = MyCli(myclirc=myclirc) + + assert cli.plot_theme == 'carbong90' + + +def test_init_configures_kitty_image_protocol(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + patch_constructor_side_effects(monkeypatch) + myclirc = write_myclirc( + tmp_path, + """ + [dataframe] + image_protocol = kitty + """, + ) + + cli = MyCli(myclirc=myclirc) + + assert cli.image_protocol == 'kitty' + + +def test_init_reports_invalid_image_protocol(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + patch_constructor_side_effects(monkeypatch) + echo_calls: list[tuple[str, dict[str, Any]]] = [] + monkeypatch.setattr(MyCli, 'echo', lambda self, message, **kwargs: echo_calls.append((message, kwargs))) + myclirc = write_myclirc( + tmp_path, + """ + [dataframe] + image_protocol = sixel + """, + ) + + cli = MyCli(myclirc=myclirc) + + assert cli.image_protocol == 'none' + assert echo_calls == [('Invalid config option provided for image_protocol (sixel); disabling.', {'err': True, 'fg': 'red'})] + + def test_init_honors_explicit_show_warnings(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: patch_constructor_side_effects(monkeypatch) show_warnings_calls: list[bool] = [] diff --git a/test/pytests/test_main_modes_repl.py b/test/pytests/test_main_modes_repl.py index 322ff56dd..d17eadebe 100644 --- a/test/pytests/test_main_modes_repl.py +++ b/test/pytests/test_main_modes_repl.py @@ -154,6 +154,10 @@ def make_repl_cli(sqlexecute: Any | None = None) -> Any: cli.null_string = '' cli.numeric_alignment = 'right' cli.binary_display = None + cli.image_protocol = 'none' + cli.plot_scale_factor = 1.0 + cli.plot_ppi = 200 + cli.plot_theme = 'carbong90' cli.prompt_session = None cli.post_redirect_command = None cli.logfile = None @@ -1178,7 +1182,7 @@ def run(self, text: str) -> Iterator[SQLResult]: cli = make_repl_cli(sqlexecute) transform = object() prepare_calls: list[tuple[str, str]] = [] - run_calls: list[object] = [] + run_calls: list[tuple[object, str, float, int, str]] = [] output_flags: list[bool] = [] hook_calls: list[tuple[str, str]] = [] @@ -1186,8 +1190,16 @@ def prepare(sql: str, expression: str) -> object: prepare_calls.append((sql, expression)) return transform - def run(received_transform: object, results: Iterator[SQLResult]) -> SQLResult: - run_calls.append(received_transform) + def run( + received_transform: object, + results: Iterator[SQLResult], + *, + image_protocol: str, + plot_scale_factor: float, + plot_ppi: int, + plot_theme: str, + ) -> SQLResult: + run_calls.append((received_transform, image_protocol, plot_scale_factor, plot_ppi, plot_theme)) assert list(results) == [SQLResult(header=['id'], rows=[(1,)])] return SQLResult(header=['count'], rows=[(1,)]) @@ -1201,13 +1213,17 @@ def run(received_transform: object, results: Iterator[SQLResult]) -> SQLResult: lambda command, filename: hook_calls.append((command, filename)), ) cli.post_redirect_command = 'post {}' + cli.image_protocol = 'iterm2' + cli.plot_scale_factor = 1.5 + cli.plot_ppi = 144 + cli.plot_theme = 'dark' command = f'SELECT * FROM orders .| df.group_by(\'customer_id\').len() {terminator}' repl_mode._one_iteration(cli, repl_mode.ReplState(), command) assert sqlexecute.calls == ['SELECT * FROM orders'] assert prepare_calls == [('SELECT * FROM orders', "df.group_by('customer_id').len()")] - assert run_calls == [transform] + assert run_calls == [(transform, 'iterm2', 1.5, 144, 'dark')] assert output_flags == [True] assert hook_calls == [] assert cli.log_queries == [command] @@ -1256,7 +1272,20 @@ def prepare(sql: str, expression: str | None) -> object: prepare_calls.append((sql, expression)) return transform - def run(received_transform: object, results: Iterator[SQLResult], path: str) -> SQLResult: + 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 image_protocol == 'none' + assert plot_scale_factor == 1.0 + assert plot_ppi == 200 + assert plot_theme == 'carbong90' assert list(results) == [SQLResult(header=['id'], rows=[(1,)])] run_calls.append((received_transform, path)) return SQLResult(status=f'Wrote 1 rows to {path}.') @@ -1306,7 +1335,20 @@ def prepare(sql: str, expression: str | None) -> object: prepare_calls.append((sql, expression)) return transform - def run(received_transform: object, results: Iterator[SQLResult], path: str) -> SQLResult: + 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 image_protocol == 'none' + assert plot_scale_factor == 1.0 + assert plot_ppi == 200 + assert plot_theme == 'carbong90' assert received_transform is transform assert list(results) == [SQLResult(header=['id'], rows=[(1,)])] run_calls.append(path) @@ -1351,7 +1393,9 @@ def run(self, text: str) -> Iterator[SQLResult]: monkeypatch.setattr( repl_mode, 'run_polars_transform', - lambda received_transform, results, path: SQLResult(status=f'Wrote 1 rows to {path}.'), + lambda received_transform, results, path, *, image_protocol, plot_scale_factor, plot_ppi, plot_theme: SQLResult( + status=f'Wrote 1 rows to {path}.' + ), ) def raise_hook_error(command: str, filename: str) -> None: @@ -1384,7 +1428,16 @@ def run(self, text: str) -> Iterator[SQLResult]: monkeypatch.setattr(repl_mode, 'prepare_polars_transform', lambda sql, expression: object()) - def fail_write(received_transform: object, results: Iterator[SQLResult], path: str) -> SQLResult: + def fail_write( + received_transform: object, + results: Iterator[SQLResult], + path: str, + *, + image_protocol: str, + plot_scale_factor: float, + plot_ppi: int, + plot_theme: str, + ) -> SQLResult: raise repl_mode.PolarsTransformError('write failed') monkeypatch.setattr(repl_mode, 'run_polars_transform', fail_write) @@ -1417,7 +1470,9 @@ def run(self, text: str) -> Iterator[SQLResult]: monkeypatch.setattr( repl_mode, 'run_polars_transform', - lambda transform, results: (_ for _ in ()).throw(repl_mode.PolarsTransformError('Polars expression failed')), + lambda transform, results, *, image_protocol, plot_scale_factor, plot_ppi, plot_theme: (_ for _ in ()).throw( + repl_mode.PolarsTransformError('Polars expression failed') + ), ) command = 'SELECT * FROM orders .| df.bad_method()' diff --git a/test/pytests/test_output.py b/test/pytests/test_output.py index b64bfb459..696e20656 100644 --- a/test/pytests/test_output.py +++ b/test/pytests/test_output.py @@ -79,6 +79,39 @@ def test_echo_logs_and_prints(monkeypatch: pytest.MonkeyPatch) -> None: assert printed == [('message', {'fg': 'red'})] +def test_output_emits_iterm2_image_without_text_sinks(monkeypatch: pytest.MonkeyPatch) -> None: + cli = make_bare_mycli() + cli.explicit_pager = False + cli.get_output_margin = lambda status=None: 1 # type: ignore[assignment] + logged: list[Any] = [] + emitted: list[str] = [] + cli.log_output = lambda value: logged.append(value) # type: ignore[assignment] + monkeypatch.setattr(click, 'echo', lambda value=None, **_kwargs: emitted.append(value)) + + OutputMixin.output(cli, itertools.chain(), SQLResult(image=b'png', image_protocol='iterm2')) + + assert logged == [] + assert emitted == ['\x1b]1337;File=name=Y2hhcnQucG5n;size=3;width=auto;height=auto;preserveAspectRatio=1;inline=1:cG5n\x07'] + + +def test_output_emits_kitty_image_in_base64_chunks(monkeypatch: pytest.MonkeyPatch) -> None: + cli = make_bare_mycli() + cli.explicit_pager = False + cli.get_output_margin = lambda status=None: 1 # type: ignore[assignment] + emitted: list[tuple[str | None, dict[str, Any]]] = [] + monkeypatch.setattr(click, 'echo', lambda value=None, **kwargs: emitted.append((value, kwargs))) + image = b'x' * 4000 + + OutputMixin.output(cli, itertools.chain(), SQLResult(image=image, image_protocol='kitty')) + + encoded = output_module.base64.b64encode(image).decode('ascii') + assert emitted == [ + (f'\x1b_Ga=T,f=100,m=1;{encoded[:4096]}\x1b\\', {'nl': False}), + (f'\x1b_Gm=0;{encoded[4096:]}\x1b\\', {'nl': False}), + (None, {}), + ] + + def test_get_output_margin_renders_prompt_once_and_counts_status_lines(monkeypatch: pytest.MonkeyPatch) -> None: cli = make_bare_mycli() cli.prompt_lines = 0 diff --git a/test/pytests/test_polars_transform.py b/test/pytests/test_polars_transform.py index 4536c485e..44375173d 100644 --- a/test/pytests/test_polars_transform.py +++ b/test/pytests/test_polars_transform.py @@ -14,7 +14,7 @@ run_polars_transform, ) from mycli.packages.sqlresult import SQLResult -from mycli.types import OutputMode +from mycli.types import ImageProtocol, OutputMode class FakeDataFrame: @@ -68,10 +68,55 @@ class FailingPolars: DataFrame = FailingDataFrame -class FakeAltair: +class FakeTopLevelMixin: pass +class FakeChart(FakeTopLevelMixin): + scale_factors: list[float] = [] + ppis: list[int] = [] + + def __init__(self, _data: FakeDataFrame) -> None: + self.save_calls: list[str] = [] + + 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') + + +class FailingChart(FakeTopLevelMixin): + def save(self, _file: Any, **_kwargs: str) -> None: + raise OSError('renderer failed') + + +class FakeTheme: + enabled: list[str] = [] + + @classmethod + def enable(cls, name: str) -> None: + cls.enabled.append(name) + + +class FakeAltair: + TopLevelMixin = FakeTopLevelMixin + theme = FakeTheme + + @staticmethod + def Chart(data: FakeDataFrame) -> FakeChart: + return FakeChart(data) + + +class FailingAltair: + TopLevelMixin = FakeTopLevelMixin + theme = FakeTheme + + @staticmethod + def Chart(_data: FakeDataFrame) -> FailingChart: + return FailingChart() + + def make_transform(expression: str) -> PolarsTransform: return PolarsTransform( sql='SELECT id FROM orders', @@ -258,6 +303,20 @@ def fail_altair_import(name: str, *args: Any, **kwargs: Any) -> Any: polars_transform._load_altair() +def test_load_vl_convert_reports_missing_dependency(monkeypatch: pytest.MonkeyPatch) -> None: + original_import = __import__ + + def fail_vl_convert_import(name: str, *args: Any, **kwargs: Any) -> Any: + if name == 'vl_convert': + raise ImportError('not installed') + return original_import(name, *args, **kwargs) + + monkeypatch.setattr('builtins.__import__', fail_vl_convert_import) + + with pytest.raises(PolarsTransformError, match='requires vl-convert-python'): + polars_transform._load_vl_convert() + + def test_run_polars_transform_materializes_and_renders_returned_dataframe() -> None: result = run_polars_transform( make_transform('df'), @@ -378,6 +437,114 @@ def test_run_polars_transform_reports_non_dataframe_result() -> None: assert result.status_plain == "Nothing could be displayed for return type: " +def test_run_polars_transform_reports_disabled_altair_chart_output() -> None: + result = run_polars_transform( + make_transform('alt.Chart(df)'), + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + assert result.status_plain == 'image_protocol is unset in ~/.myclirc. Inline plotting is disabled.' + + +@pytest.mark.parametrize('image_protocol', ['iterm2', 'kitty']) +def test_run_polars_transform_renders_altair_chart_for_image_protocol( + monkeypatch: pytest.MonkeyPatch, + image_protocol: ImageProtocol, +) -> None: + load_calls: list[None] = [] + FakeChart.scale_factors = [] + FakeChart.ppis = [] + 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,)])]), + image_protocol=image_protocol, + plot_scale_factor=1.5, + plot_ppi=144, + plot_theme='dark', + ) + + assert load_calls == [None] + assert FakeChart.scale_factors == [1.5] + assert FakeChart.ppis == [144] + assert FakeTheme.enabled == ['dark'] + assert result == SQLResult(image=b'png image', image_protocol=image_protocol) + + +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) + + run_polars_transform( + make_transform('alt.Chart(df)'), + iter([SQLResult(header=['id'], rows=[(1,)])]), + image_protocol='kitty', + ) + + assert FakeChart.ppis == [200] + + +def test_run_polars_transform_reports_invalid_altair_theme(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(polars_transform, '_load_vl_convert', lambda: None) + + def fail_enable(_name: str) -> None: + raise ValueError('unknown theme') + + monkeypatch.setattr(FakeTheme, 'enable', fail_enable) + + 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)'), + 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: + 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 render Altair chart: OSError: renderer failed'): + run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,)])]), + image_protocol='iterm2', + ) + + +def test_run_polars_transform_reports_missing_altair_renderer(monkeypatch: pytest.MonkeyPatch) -> None: + def fail_load() -> None: + raise PolarsTransformError('Altair plot rendering requires vl-convert-python. Install mycli[dataframe].') + + monkeypatch.setattr(polars_transform, '_load_vl_convert', fail_load) + + with pytest.raises(PolarsTransformError, match='requires vl-convert-python'): + run_polars_transform( + make_transform('alt.Chart(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='must return a DataFrame or Series'): + run_polars_transform( + make_transform('alt.Chart(df)'), + iter([SQLResult(header=['id'], rows=[(1,)])]), + 'chart.parquet', + image_protocol='iterm2', + ) + + def test_prepare_polars_transform_supports_direct_parquet_output(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr('mycli.packages.polars_transform._load_polars', lambda: FakePolars)