Column picker for copy & export

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)

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

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

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:
         ...
 

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)

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