From 31f829239f689c2b7c12f9253cf63ca721b3a0b7 Mon Sep 17 00:00:00 2001 From: Peter Adams <18162810+Maxteabag@users.noreply.github.com> Date: Mon, 25 May 2026 16:08:24 +0200 Subject: [PATCH] feat(results): column picker for copy + export MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a "Columns…" entry in the results yank menu and inside Copy as… and Export… submenus. The picker shows a checklist of result columns — toggle with space, a = all, n = none — and applies the selected subset to a quick TSV copy, a chosen Copy-as format, or an export to file. --- sqlit/core/keymap.py | 5 + sqlit/domains/results/formatters.py | 10 ++ sqlit/domains/results/ui/mixins/results.py | 103 ++++++++++++++++++- sqlit/shared/ui/protocols/results.py | 13 +++ sqlit/shared/ui/screens/column_picker.py | 114 +++++++++++++++++++++ tests/test_result_formatters.py | 33 ++++++ 6 files changed, 274 insertions(+), 4 deletions(-) create mode 100644 sqlit/shared/ui/screens/column_picker.py diff --git a/sqlit/core/keymap.py b/sqlit/core/keymap.py index 23a693d5..7f320e49 100644 --- a/sqlit/core/keymap.py +++ b/sqlit/core/keymap.py @@ -295,12 +295,14 @@ def _build_leader_commands(self) -> list[LeaderCommandDef]: LeaderCommandDef("c", "cell", "Copy cell", "Copy", menu="ry"), LeaderCommandDef("y", "row", "Copy row", "Copy", menu="ry"), LeaderCommandDef("a", "all", "Copy all", "Copy", menu="ry"), + LeaderCommandDef("o", "columns", "Copy columns...", "Copy", menu="ry"), LeaderCommandDef("f", "format", "Copy as...", "Copy", menu="ry"), LeaderCommandDef("e", "export", "Export...", "Export", menu="ry"), # rye results export menu LeaderCommandDef("c", "csv", "Export as CSV", "Export", menu="rye"), LeaderCommandDef("j", "json", "Export as JSON", "Export", menu="rye"), LeaderCommandDef("m", "markdown", "Export as Markdown", "Export", menu="rye"), + LeaderCommandDef("o", "columns", "Columns...", "Export", menu="rye"), # ryf copy-as format picker LeaderCommandDef("m", "markdown", "Copy as Markdown...", "Copy as", menu="ryf"), LeaderCommandDef("j", "json", "Copy as JSON...", "Copy as", menu="ryf"), @@ -310,14 +312,17 @@ def _build_leader_commands(self) -> list[LeaderCommandDef]: LeaderCommandDef("c", "cell", "Cell as Markdown", "Copy as", menu="ryfm"), LeaderCommandDef("y", "row", "Row as Markdown", "Copy as", menu="ryfm"), LeaderCommandDef("a", "all", "All as Markdown", "Copy as", menu="ryfm"), + LeaderCommandDef("o", "columns", "Columns as Markdown...", "Copy as", menu="ryfm"), # ryfj json scope LeaderCommandDef("c", "cell", "Cell as JSON", "Copy as", menu="ryfj"), LeaderCommandDef("y", "row", "Row as JSON", "Copy as", menu="ryfj"), LeaderCommandDef("a", "all", "All as JSON", "Copy as", menu="ryfj"), + LeaderCommandDef("o", "columns", "Columns as JSON...", "Copy as", menu="ryfj"), # ryfc csv scope LeaderCommandDef("c", "cell", "Cell as CSV", "Copy as", menu="ryfc"), LeaderCommandDef("y", "row", "Row as CSV", "Copy as", menu="ryfc"), LeaderCommandDef("a", "all", "All as CSV", "Copy as", menu="ryfc"), + LeaderCommandDef("o", "columns", "Columns as CSV...", "Copy as", menu="ryfc"), # rg results g motion menu (vim-style gg) LeaderCommandDef("g", "first_row", "Go to first row", "Go to", menu="rg"), # vy value view yank menu (tree mode) diff --git a/sqlit/domains/results/formatters.py b/sqlit/domains/results/formatters.py index 82b7948e..804035c4 100644 --- a/sqlit/domains/results/formatters.py +++ b/sqlit/domains/results/formatters.py @@ -94,3 +94,13 @@ class ResultFormat: "json": ResultFormat("json", "JSON", "json", format_json), "markdown": ResultFormat("markdown", "Markdown", "md", format_markdown), } + + +def project_columns( + columns: Columns, rows: Rows, indices: Sequence[int] +) -> tuple[list[str], list[tuple[Any, ...]]]: + """Return (columns, rows) restricted to the given column indices.""" + ordered = [i for i in indices if 0 <= i < len(columns)] + new_cols = [columns[i] for i in ordered] + new_rows = [tuple(row[i] for i in ordered if i < len(row)) for row in rows] + return new_cols, new_rows diff --git a/sqlit/domains/results/ui/mixins/results.py b/sqlit/domains/results/ui/mixins/results.py index 89300172..e84becf8 100644 --- a/sqlit/domains/results/ui/mixins/results.py +++ b/sqlit/domains/results/ui/mixins/results.py @@ -35,6 +35,7 @@ class ResultsMixin: _last_result_columns: list[str] = [] _last_result_rows: list[tuple[Any, ...]] = [] + _export_column_indices: list[int] | None = None _last_result_row_count: int = 0 _tooltip_cell_coord: tuple[int, int] | None = None _tooltip_showing: bool = False @@ -603,6 +604,24 @@ def action_ry_all(self: ResultsMixinHost) -> None: if table: self._flash_table_yank(table, "all") + def action_ry_columns(self: ResultsMixinHost) -> None: + """Pick a column subset, then copy all rows of those columns as TSV.""" + self._clear_leader_pending() + table, columns, rows, _stacked = self._get_active_results_context() + if not table or table.row_count <= 0 or not columns: + self.notify("No results", severity="warning") + return + + def do_copy(indices: list[int]) -> None: + from sqlit.domains.results.formatters import project_columns + + sub_cols, sub_rows = project_columns(columns, rows, indices) + text = self._format_tsv(sub_cols, sub_rows) + self._copy_text(text) + self._flash_table_yank(table, "all") + + self._pick_columns(columns, on_confirm=do_copy) + def action_ry_export(self: ResultsMixinHost) -> None: """Open the export submenu.""" self._clear_leader_pending() @@ -642,6 +661,9 @@ def _show_export_dialog(self: ResultsMixinHost, fmt_key: str) -> None: def handle_result(filename: str | None) -> None: if filename: self._save_export_file(filename, fmt_key) + else: + # Cancelled: don't leak a column subset to the next export attempt. + self._export_column_indices = None self.push_screen( FilePickerScreen( @@ -653,22 +675,33 @@ def handle_result(filename: str | None) -> None: ) def _save_export_file(self: ResultsMixinHost, filename: str, fmt_key: str) -> None: - """Save the export file to disk.""" + """Save the export file to disk. + + Honors a one-shot column subset set via the `rye o` (Columns…) flow; + the subset is consumed and cleared once the export runs. + """ from pathlib import Path - from sqlit.domains.results.formatters import FORMATS + from sqlit.domains.results.formatters import FORMATS, project_columns try: fmt = FORMATS[fmt_key] - content = fmt.formatter(self._last_result_columns, self._last_result_rows) + cols = list(self._last_result_columns) + rows = list(self._last_result_rows) + subset = getattr(self, "_export_column_indices", None) + if subset: + cols, rows = project_columns(cols, rows, subset) + content = fmt.formatter(cols, rows) path = Path(filename).expanduser() path.write_text(content, encoding="utf-8") - row_count = len(self._last_result_rows) + row_count = len(rows) self.notify(f"Saved {row_count} rows to {path.name}") except Exception as e: self.notify(f"Failed to save: {e}", severity="error") + finally: + self._export_column_indices = None # ------------------------------------------------------------------ # Copy-as: ryf menu — pick format, then scope (cell/row/all). @@ -727,6 +760,35 @@ def action_ryfc_row(self: ResultsMixinHost) -> None: def action_ryfc_all(self: ResultsMixinHost) -> None: self._copy_scope_as_format("csv", "all") + def action_ryfm_columns(self: ResultsMixinHost) -> None: + self._copy_columns_as_format("markdown") + + def action_ryfj_columns(self: ResultsMixinHost) -> None: + self._copy_columns_as_format("json") + + def action_ryfc_columns(self: ResultsMixinHost) -> None: + self._copy_columns_as_format("csv") + + def action_rye_columns(self: ResultsMixinHost) -> None: + """Pick columns, then pick a format to export with.""" + self._clear_leader_pending() + if not self._last_result_columns or not self._last_result_rows: + self.notify("No results to export", severity="warning") + return + self._pick_columns( + self._last_result_columns, + on_confirm=lambda indices: self._start_leader_pending_for_export_with_columns( + indices + ), + ) + + def _start_leader_pending_for_export_with_columns( + self: ResultsMixinHost, indices: list[int] + ) -> None: + """After a column pick, store the subset and open the format submenu.""" + self._export_column_indices = indices + self._start_leader_pending("rye") + def _copy_scope_as_format( self: ResultsMixinHost, fmt_key: str, scope: str ) -> None: @@ -765,6 +827,39 @@ def _copy_scope_as_format( self._copy_text(content) self._flash_table_yank(table, scope) + def _copy_columns_as_format(self: ResultsMixinHost, fmt_key: str) -> None: + """Open the column picker, then copy the subset as fmt_key.""" + self._clear_leader_pending() + table, columns, rows, _stacked = self._get_active_results_context() + if not table or table.row_count <= 0 or not columns: + self.notify("No results", severity="warning") + return + + def do_copy(indices: list[int]) -> None: + from sqlit.domains.results.formatters import FORMATS, project_columns + + sub_cols, sub_rows = project_columns(columns, rows, indices) + content = FORMATS[fmt_key].formatter(sub_cols, sub_rows) + self._copy_text(content) + self._flash_table_yank(table, "all") + + self._pick_columns(columns, on_confirm=do_copy) + + def _pick_columns( + self: ResultsMixinHost, + columns: list[str], + *, + on_confirm: Any, + ) -> None: + """Show the column-picker modal; call on_confirm(indices) when confirmed.""" + from sqlit.shared.ui.screens.column_picker import ColumnPickerScreen + + def handle(result: list[int] | None) -> None: + if result: + on_confirm(result) + + self.push_screen(ColumnPickerScreen(columns), handle) + def _copy_column_values(self: ResultsMixinHost) -> None: """Copy every value in the focused column as a SQL-ready list.""" from sqlit.domains.results.formatters import format_values_list diff --git a/sqlit/shared/ui/protocols/results.py b/sqlit/shared/ui/protocols/results.py index 5b8fcb82..66c71985 100644 --- a/sqlit/shared/ui/protocols/results.py +++ b/sqlit/shared/ui/protocols/results.py @@ -124,6 +124,19 @@ def _copy_scope_as_format(self, fmt_key: str, scope: str) -> None: def _copy_column_values(self) -> None: ... + def _copy_columns_as_format(self, fmt_key: str) -> None: + ... + + def _pick_columns(self, columns: list[str], *, on_confirm: Any) -> None: + ... + + def _start_leader_pending_for_export_with_columns( + self, indices: list[int] + ) -> None: + ... + + _export_column_indices: list[int] | None + def _show_single_result_mode(self) -> None: ... diff --git a/sqlit/shared/ui/screens/column_picker.py b/sqlit/shared/ui/screens/column_picker.py new file mode 100644 index 00000000..6ac23774 --- /dev/null +++ b/sqlit/shared/ui/screens/column_picker.py @@ -0,0 +1,114 @@ +"""Column-picker modal used by copy/export actions.""" + +from __future__ import annotations + +from textual.app import ComposeResult +from textual.binding import Binding +from textual.screen import ModalScreen +from textual.widgets import SelectionList +from textual.widgets.selection_list import Selection + +from sqlit.shared.ui.widgets import Dialog + + +class _ColumnSelectionList(SelectionList): + """SelectionList that releases Enter so the parent screen can confirm. + + Upstream binds Enter to toggle (alongside Space). The picker needs Enter + to mean "Confirm", so we override BINDINGS to keep only Space for toggle. + """ + + BINDINGS = [Binding("space", "select")] + + +class ColumnPickerScreen(ModalScreen[list[int] | None]): + """Pick which columns to include in a copy/export operation. + + Returns the list of selected column indices (preserving original order), + or None on cancel. + """ + + BINDINGS = [ + Binding("enter", "confirm", "Confirm", show=False, priority=True), + Binding("escape", "cancel", "Cancel", show=False, priority=True), + Binding("a", "select_all", "Select all", show=False, priority=True), + Binding("n", "deselect_all", "Deselect all", show=False, priority=True), + ] + + CSS = """ + ColumnPickerScreen { + align: center middle; + background: transparent; + } + + #column-picker-dialog { + width: 70%; + max-width: 100; + min-width: 60; + max-height: 80%; + } + + #column-list { + height: auto; + max-height: 20; + background: $surface; + } + """ + + def __init__( + self, + columns: list[str], + initial_selected: list[int] | None = None, + *, + title: str = "Select columns", + ) -> None: + super().__init__() + self._columns = list(columns) + if initial_selected is None: + self._initial = set(range(len(columns))) + else: + self._initial = set(initial_selected) + self._title = title + + def compose(self) -> ComposeResult: + selections = [ + Selection(name, idx, idx in self._initial) + for idx, name in enumerate(self._columns) + ] + shortcuts = [ + ("Confirm", "enter"), + ("Toggle", "space"), + ("All", "a"), + ("None", "n"), + ("Cancel", "esc"), + ] + with Dialog(id="column-picker-dialog", title=self._title, shortcuts=shortcuts): + yield _ColumnSelectionList(*selections, id="column-list") + + def on_mount(self) -> None: + self.query_one("#column-list", SelectionList).focus() + + def _selected_indices(self) -> list[int]: + sel_list = self.query_one("#column-list", SelectionList) + return sorted(int(v) for v in sel_list.selected) + + def action_confirm(self) -> None: + indices = self._selected_indices() + if not indices: + self.app.notify("Select at least one column", severity="warning") + return + self.dismiss(indices) + + def action_cancel(self) -> None: + self.dismiss(None) + + def action_select_all(self) -> None: + self.query_one("#column-list", SelectionList).select_all() + + def action_deselect_all(self) -> None: + self.query_one("#column-list", SelectionList).deselect_all() + + def check_action(self, action: str, parameters: tuple) -> bool | None: + if self.app.screen is not self: + return False + return super().check_action(action, parameters) diff --git a/tests/test_result_formatters.py b/tests/test_result_formatters.py index 8d5e1627..ca70cecb 100644 --- a/tests/test_result_formatters.py +++ b/tests/test_result_formatters.py @@ -12,6 +12,7 @@ format_json, format_markdown, format_values_list, + project_columns, ) @@ -86,3 +87,35 @@ def test_format_registry_keys_and_extensions(): def test_each_format_runs_on_sample(key): out = FORMATS[key].formatter(COLS, ROWS) assert isinstance(out, str) and out + + +def test_project_columns_subset(): + cols, rows = project_columns(COLS, ROWS, [0, 2]) + assert cols == ["id", "note"] + assert rows == [(1, "a|b"), (2, None), (3, "x")] + + +def test_project_columns_reorders_to_given_indices_sorted(): + # The action layer passes already-sorted indices; project_columns honors order. + cols, rows = project_columns(COLS, ROWS, [2, 0]) + assert cols == ["note", "id"] + assert rows[0] == ("a|b", 1) + + +def test_project_columns_ignores_out_of_range(): + cols, _rows = project_columns(COLS, ROWS, [0, 99]) + assert cols == ["id"] + + +def test_project_columns_empty_indices_yields_empty_rows(): + cols, rows = project_columns(COLS, ROWS, []) + assert cols == [] + assert rows == [(), (), ()] + + +def test_project_columns_composes_with_csv(): + cols, rows = project_columns(COLS, ROWS, [1]) + out = format_csv(cols, rows) + # csv.writer quotes embedded newlines so the multi-line cell stays intact. + assert out.startswith("name\r\nAlice\r\nBob\r\n") + assert '"C\nD"' in out