Merge pull request #113 from Maxteabag/feature/vim-count-prefixf

feat: add vim count prefix support for motions and operators

Author: Maxteabag

sqlit/core/input_context.py

@@ -33,3 +33,4 @@ class InputContext:
     last_result_is_error: bool
     has_results: bool
     stacked_result_count: int = 0
+    count_buffer: str = ""

sqlit/domains/query/ui/mixins/query_editing_cursor.py

@@ -9,26 +9,46 @@ class QueryEditingCursorMixin:
     """Cursor movement and navigation for the query editor."""
 
     def _move_with_motion(self: QueryMixinHost, motion_key: str, char: str | None = None) -> None:
-        """Move cursor using a vim motion."""
+        """Move cursor using a vim motion, with optional count prefix support."""
         from sqlit.domains.query.editing import MOTIONS
 
         motion_func = MOTIONS.get(motion_key)
         if not motion_func:
             return
 
+        # Get count prefix (if any)
+        count = self._get_and_clear_count() or 1
+
         text = self.query_input.text
         row, col = self.query_input.cursor_location
-        result = motion_func(text, row, col, char)
-        self.query_input.cursor_location = (result.position.row, result.position.col)
+
+        # Apply motion `count` times
+        for _ in range(count):
+            result = motion_func(text, row, col, char)
+            new_row, new_col = result.position.row, result.position.col
+            # Stop if motion didn't move (hit boundary)
+            if (new_row, new_col) == (row, col):
+                break
+            row, col = new_row, new_col
+
+        self.query_input.cursor_location = (row, col)
 
     def action_g_leader_key(self: QueryMixinHost) -> None:
         """Show the g motion leader menu."""
         self._start_leader_pending("g")
 
     def action_g_first_line(self: QueryMixinHost) -> None:
-        """Go to first line (gg)."""
+        """Go to first line (gg), or to line N with count prefix (e.g., 3gg)."""
         self._clear_leader_pending()
-        self.query_input.cursor_location = (0, 0)
+        count = self._get_and_clear_count()
+        if count is not None:
+            lines = self.query_input.text.split("\n")
+            num_lines = len(lines)
+            target_row = min(count - 1, num_lines - 1)
+            target_row = max(0, target_row)
+            self.query_input.cursor_location = (target_row, 0)
+        else:
+            self.query_input.cursor_location = (0, 0)
 
     def action_g_word_end_back(self: QueryMixinHost) -> None:
         """Go to end of previous word (ge)."""
@@ -56,32 +76,20 @@ def action_g_execute_query_atomic(self: QueryMixinHost) -> None:
         self.action_execute_query_atomic()
 
     def action_cursor_left(self: QueryMixinHost) -> None:
-        """Move cursor left (h in normal mode)."""
-        row, col = self.query_input.cursor_location
-        self.query_input.cursor_location = (row, max(0, col - 1))
+        """Move cursor left (h in normal mode), with count support."""
+        self._move_with_motion("h")
 
     def action_cursor_right(self: QueryMixinHost) -> None:
-        """Move cursor right (l in normal mode)."""
-        lines = self.query_input.text.split("\n")
-        row, col = self.query_input.cursor_location
-        line_len = len(lines[row]) if row < len(lines) else 0
-        self.query_input.cursor_location = (row, min(col + 1, line_len))
+        """Move cursor right (l in normal mode), with count support."""
+        self._move_with_motion("l")
 
     def action_cursor_up(self: QueryMixinHost) -> None:
-        """Move cursor up (k in normal mode)."""
-        lines = self.query_input.text.split("\n")
-        row, col = self.query_input.cursor_location
-        new_row = max(0, row - 1)
-        new_col = min(col, len(lines[new_row]) if new_row < len(lines) else 0)
-        self.query_input.cursor_location = (new_row, new_col)
+        """Move cursor up (k in normal mode), with count support."""
+        self._move_with_motion("k")
 
     def action_cursor_down(self: QueryMixinHost) -> None:
-        """Move cursor down (j in normal mode)."""
-        lines = self.query_input.text.split("\n")
-        row, col = self.query_input.cursor_location
-        new_row = min(row + 1, len(lines) - 1)
-        new_col = min(col, len(lines[new_row]) if new_row < len(lines) else 0)
-        self.query_input.cursor_location = (new_row, new_col)
+        """Move cursor down (j in normal mode), with count support."""
+        self._move_with_motion("j")
 
     def action_cursor_word_forward(self: QueryMixinHost) -> None:
         """Move cursor to next word (w)."""
@@ -108,8 +116,18 @@ def action_cursor_line_end(self: QueryMixinHost) -> None:
         self._move_with_motion("$")
 
     def action_cursor_last_line(self: QueryMixinHost) -> None:
-        """Move cursor to last line (G)."""
-        self._move_with_motion("G")
+        """Move cursor to last line (G), or to line N with count prefix (e.g., 25G)."""
+        count = self._get_and_clear_count()
+        if count is not None:
+            # Go to specific line (1-indexed)
+            lines = self.query_input.text.split("\n")
+            num_lines = len(lines)
+            target_row = min(count - 1, num_lines - 1)  # Convert to 0-indexed, clamp
+            target_row = max(0, target_row)
+            self.query_input.cursor_location = (target_row, 0)
+        else:
+            # Go to last line
+            self._move_with_motion("G")
 
     def action_cursor_matching_bracket(self: QueryMixinHost) -> None:
         """Move cursor to matching bracket (%)."""

sqlit/domains/query/ui/mixins/query_editing_operators.py

@@ -10,22 +10,14 @@ class QueryEditingOperatorsMixin:
     """Delete/yank/change operator actions for the query editor."""
 
     def action_delete_line(self: QueryMixinHost) -> None:
-        """Delete the current line in the query editor."""
+        """Delete the current line (dd), with count support for multi-line delete."""
         self._clear_leader_pending()
-        result = edit_delete.delete_line(
-            self.query_input.text,
-            *self.query_input.cursor_location,
-        )
-        self._apply_edit_result(result)
+        self._delete_with_motion("_")  # _ is the current line motion
 
     def action_delete_word(self: QueryMixinHost) -> None:
-        """Delete forward word starting at cursor."""
+        """Delete forward word (dw), with count support."""
         self._clear_leader_pending()
-        result = edit_delete.delete_word(
-            self.query_input.text,
-            *self.query_input.cursor_location,
-        )
-        self._apply_edit_result(result)
+        self._delete_with_motion("w")
 
     def action_delete_word_back(self: QueryMixinHost) -> None:
         """Delete word backwards from cursor."""
@@ -200,24 +192,57 @@ def handle_result(obj_char: str | None) -> None:
         self.push_screen(TextObjectMenuScreen(mode, operator="delete"), handle_result)
 
     def _delete_with_motion(self: QueryMixinHost, motion_key: str, char: str | None = None) -> None:
-        """Execute delete with a motion."""
-        from sqlit.domains.query.editing import MOTIONS, operator_delete
+        """Execute delete with a motion, with optional count prefix support."""
+        from sqlit.domains.query.editing import MOTIONS, MotionType, Position, Range, operator_delete
 
         motion_func = MOTIONS.get(motion_key)
         if not motion_func:
             return
 
+        # Get count prefix (if any)
+        count = self._get_and_clear_count() or 1
+
         text = self.query_input.text
         row, col = self.query_input.cursor_location
+        lines = text.split("\n")
 
         result = motion_func(text, row, col, char)
         if not result.range:
             return
 
+        final_range = result.range
+
+        # Handle count for line motions (e.g., 3dd deletes 3 lines)
+        if motion_key == "_" and count > 1:
+            # _ is the current line motion; expand to cover `count` lines
+            start_row = row
+            end_row = min(row + count - 1, len(lines) - 1)
+            end_col = len(lines[end_row]) if end_row < len(lines) else 0
+            final_range = Range(
+                Position(start_row, 0),
+                Position(end_row, end_col),
+                MotionType.LINEWISE,
+            )
+        elif count > 1:
+            # For other motions, iterate to expand range
+            end_row, end_col = result.position.row, result.position.col
+            for _ in range(count - 1):
+                next_result = motion_func(text, end_row, end_col, char)
+                if (next_result.position.row, next_result.position.col) == (end_row, end_col):
+                    break  # Motion didn't move
+                end_row, end_col = next_result.position.row, next_result.position.col
+            # Rebuild range from original position to final position
+            final_range = Range(
+                result.range.start,
+                Position(end_row, end_col),
+                result.range.motion_type,
+                result.range.inclusive,
+            )
+
         # Push undo state before delete
         self._push_undo_state()
 
-        op_result = operator_delete(text, result.range)
+        op_result = operator_delete(text, final_range)
         self.query_input.text = op_result.text
         self.query_input.cursor_location = (op_result.row, op_result.col)
 
@@ -389,27 +414,58 @@ def handle_result(obj_char: str | None) -> None:
         self.push_screen(TextObjectMenuScreen(mode, operator="yank"), handle_result)
 
     def _yank_with_motion(self: QueryMixinHost, motion_key: str, char: str | None = None) -> None:
-        """Execute yank with a motion."""
-        from sqlit.domains.query.editing import MOTIONS, operator_yank
+        """Execute yank with a motion, with optional count prefix support."""
+        from sqlit.domains.query.editing import MOTIONS, MotionType, Position, Range, operator_yank
 
         motion_func = MOTIONS.get(motion_key)
         if not motion_func:
             return
 
+        # Get count prefix (if any)
+        count = self._get_and_clear_count() or 1
+
         text = self.query_input.text
         row, col = self.query_input.cursor_location
+        lines = text.split("\n")
 
         result = motion_func(text, row, col, char)
         if not result.range:
             return
 
-        op_result = operator_yank(text, result.range)
+        final_range = result.range
+
+        # Handle count for line motions (e.g., 3yy yanks 3 lines)
+        if motion_key == "_" and count > 1:
+            start_row = row
+            end_row = min(row + count - 1, len(lines) - 1)
+            end_col = len(lines[end_row]) if end_row < len(lines) else 0
+            final_range = Range(
+                Position(start_row, 0),
+                Position(end_row, end_col),
+                MotionType.LINEWISE,
+            )
+        elif count > 1:
+            # For other motions, iterate to expand range
+            end_row, end_col = result.position.row, result.position.col
+            for _ in range(count - 1):
+                next_result = motion_func(text, end_row, end_col, char)
+                if (next_result.position.row, next_result.position.col) == (end_row, end_col):
+                    break
+                end_row, end_col = next_result.position.row, next_result.position.col
+            final_range = Range(
+                result.range.start,
+                Position(end_row, end_col),
+                result.range.motion_type,
+                result.range.inclusive,
+            )
+
+        op_result = operator_yank(text, final_range)
 
         # Copy yanked text to system clipboard
         if op_result.yanked:
             self._copy_text(op_result.yanked)
             # Flash the yanked range
-            ordered = result.range.ordered()
+            ordered = final_range.ordered()
             self._flash_yank_range(
                 ordered.start.row, ordered.start.col,
                 ordered.end.row, ordered.end.col,
@@ -590,24 +646,55 @@ def handle_result(obj_char: str | None) -> None:
         self.push_screen(TextObjectMenuScreen(mode, operator="change"), handle_result)
 
     def _change_with_motion(self: QueryMixinHost, motion_key: str, char: str | None = None) -> None:
-        """Execute change with a motion (delete + enter insert mode)."""
-        from sqlit.domains.query.editing import MOTIONS, operator_change
+        """Execute change with a motion, with optional count prefix support."""
+        from sqlit.domains.query.editing import MOTIONS, MotionType, Position, Range, operator_change
 
         motion_func = MOTIONS.get(motion_key)
         if not motion_func:
             return
 
+        # Get count prefix (if any)
+        count = self._get_and_clear_count() or 1
+
         text = self.query_input.text
         row, col = self.query_input.cursor_location
+        lines = text.split("\n")
 
         result = motion_func(text, row, col, char)
         if not result.range:
             return
 
+        final_range = result.range
+
+        # Handle count for line motions (e.g., 3cc changes 3 lines)
+        if motion_key == "_" and count > 1:
+            start_row = row
+            end_row = min(row + count - 1, len(lines) - 1)
+            end_col = len(lines[end_row]) if end_row < len(lines) else 0
+            final_range = Range(
+                Position(start_row, 0),
+                Position(end_row, end_col),
+                MotionType.LINEWISE,
+            )
+        elif count > 1:
+            # For other motions, iterate to expand range
+            end_row, end_col = result.position.row, result.position.col
+            for _ in range(count - 1):
+                next_result = motion_func(text, end_row, end_col, char)
+                if (next_result.position.row, next_result.position.col) == (end_row, end_col):
+                    break
+                end_row, end_col = next_result.position.row, next_result.position.col
+            final_range = Range(
+                result.range.start,
+                Position(end_row, end_col),
+                result.range.motion_type,
+                result.range.inclusive,
+            )
+
         # Push undo state before change
         self._push_undo_state()
 
-        op_result = operator_change(text, result.range)
+        op_result = operator_change(text, final_range)
         self.query_input.text = op_result.text
         self.query_input.cursor_location = (op_result.row, op_result.col)