Implement vi mode and basic motions in pyrepl

Author: lkollar

Lib/_pyrepl/commands.py

@@ -23,6 +23,8 @@
 import os
 import time
 
+from .types import ViMode
+
 # Categories of actions:
 #  killing
 #  yanking
@@ -325,10 +327,19 @@ def do(self) -> None:
         b = r.buffer
         for _ in range(r.get_arg()):
             p = r.pos + 1
-            if p <= len(b):
-                r.pos = p
+            # In vi normal mode, don't move past the last character
+            if r.vi_mode == ViMode.NORMAL:
+                eol_pos = r.eol()
+                max_pos = max(r.bol(), eol_pos - 1) if eol_pos > r.bol() else r.bol()
+                if p <= max_pos:
+                    r.pos = p
+                else:
+                    self.reader.error("end of line")
             else:
-                self.reader.error("end of buffer")
+                if p <= len(b):
+                    r.pos = p
+                else:
+                    self.reader.error("end of buffer")
 
 
 class beginning_of_line(MotionCommand):
@@ -338,7 +349,14 @@ def do(self) -> None:
 
 class end_of_line(MotionCommand):
     def do(self) -> None:
-        self.reader.pos = self.reader.eol()
+        r = self.reader
+        eol_pos = r.eol()
+        if r.vi_mode == ViMode.NORMAL:
+            bol_pos = r.bol()
+            # Don't go past the last character (but stay at bol if line is empty)
+            r.pos = max(bol_pos, eol_pos - 1) if eol_pos > bol_pos else bol_pos
+        else:
+            r.pos = eol_pos
 
 
 class home(MotionCommand):

Lib/_pyrepl/historical_reader.py

@@ -257,19 +257,27 @@ def __post_init__(self) -> None:
         )
 
     def collect_keymap(self) -> tuple[tuple[KeySpec, CommandName], ...]:
-        return super().collect_keymap() + (
+        bindings: list[tuple[KeySpec, CommandName]] = [
             (r"\C-n", "next-history"),
             (r"\C-p", "previous-history"),
             (r"\C-o", "operate-and-get-next"),
             (r"\C-r", "reverse-history-isearch"),
             (r"\C-s", "forward-history-isearch"),
-            (r"\M-r", "restore-history"),
-            (r"\M-.", "yank-arg"),
             (r"\<page down>", "history-search-forward"),
-            (r"\x1b[6~", "history-search-forward"),
             (r"\<page up>", "history-search-backward"),
-            (r"\x1b[5~", "history-search-backward"),
-        )
+        ]
+
+        if not self.use_vi_mode:
+            bindings.extend(
+                [
+                    (r"\M-r", "restore-history"),
+                    (r"\M-.", "yank-arg"),
+                    (r"\x1b[6~", "history-search-forward"),
+                    (r"\x1b[5~", "history-search-backward"),
+                ]
+            )
+
+        return super().collect_keymap() + tuple(bindings)
 
     def select_item(self, i: int) -> None:
         self.transient_history[self.historyi] = self.get_unicode()

Lib/_pyrepl/reader.py

@@ -21,6 +21,7 @@
 
 from __future__ import annotations
 
+import itertools
 import sys
 import _colorize
 
@@ -30,6 +31,7 @@
 from . import commands, console, input
 from .utils import wlen, unbracket, disp_str, gen_colors, THEME
 from .trace import trace
+from . import vi_commands
 
 
 # types
@@ -41,6 +43,9 @@
 SYNTAX_WHITESPACE, SYNTAX_WORD, SYNTAX_SYMBOL = range(3)
 
 
+from .types import ViMode
+
+
 def make_default_syntax_table() -> dict[str, int]:
     # XXX perhaps should use some unicodedata here?
     st: dict[str, int] = {}
@@ -54,10 +59,11 @@ def make_default_syntax_table() -> dict[str, int]:
 
 def make_default_commands() -> dict[CommandName, type[Command]]:
     result: dict[CommandName, type[Command]] = {}
-    for v in vars(commands).values():
-        if isinstance(v, type) and issubclass(v, Command) and v.__name__[0].islower():
-            result[v.__name__] = v
-            result[v.__name__.replace("_", "-")] = v
+    all_commands = itertools.chain(vars(commands).values(), vars(vi_commands).values())
+    for cmd in all_commands:
+        if isinstance(cmd, type) and issubclass(cmd, Command) and cmd.__name__[0].islower():
+            result[cmd.__name__] = cmd
+            result[cmd.__name__.replace("_", "-")] = cmd
     return result
 
 
@@ -131,6 +137,67 @@ def make_default_commands() -> dict[CommandName, type[Command]]:
 )
 
 
+vi_insert_keymap: tuple[tuple[KeySpec, CommandName], ...] = tuple(
+    [binding for binding in default_keymap if not binding[0].startswith((r"\M-", r"\x1b", r"\EOF", r"\EOH"))] +
+    [(r"\<escape>", "vi-normal-mode")]
+)
+
+
+vi_normal_keymap: tuple[tuple[KeySpec, CommandName], ...] = tuple(
+    [
+        # Basic motions
+        (r"h", "left"),
+        (r"j", "down"),
+        (r"k", "up"),
+        (r"l", "right"),
+        (r"0", "beginning-of-line"),
+        (r"$", "end-of-line"),
+        (r"w", "vi-forward-word"),
+        (r"b", "backward-word"),
+        (r"e", "end-of-word"),
+        (r"^", "first-non-whitespace-character"),
+
+        # Edit commands
+        (r"x", "delete"),
+        (r"i", "vi-insert-mode"),
+        (r"a", "vi-append-mode"),
+        (r"A", "vi-append-eol"),
+        (r"I", "vi-insert-bol"),
+        (r"o", "vi-open-below"),
+        (r"O", "vi-open-above"),
+
+        # Special keys still work in normal mode
+        (r"\<left>", "left"),
+        (r"\<right>", "right"),
+        (r"\<up>", "up"),
+        (r"\<down>", "down"),
+        (r"\<home>", "beginning-of-line"),
+        (r"\<end>", "end-of-line"),
+        (r"\<delete>", "delete"),
+        (r"\<backspace>", "left"),
+
+        # Control keys (important ones that work in both modes)
+        (r"\C-c", "interrupt"),
+        (r"\C-d", "delete"),
+        (r"\C-l", "clear-screen"),
+        (r"\C-r", "reverse-history-isearch"),
+
+        # Digit args for counts (1-9, not 0 which is BOL)
+        (r"1", "digit-arg"),
+        (r"2", "digit-arg"),
+        (r"3", "digit-arg"),
+        (r"4", "digit-arg"),
+        (r"5", "digit-arg"),
+        (r"6", "digit-arg"),
+        (r"7", "digit-arg"),
+        (r"8", "digit-arg"),
+        (r"9", "digit-arg"),
+
+        (r"\<escape>", "invalid-key"),
+    ]
+)
+
+
 @dataclass(slots=True)
 class Reader:
     """The Reader class implements the bare bones of a command reader,
@@ -214,6 +281,8 @@ class Reader:
     scheduled_commands: list[str] = field(default_factory=list)
     can_colorize: bool = False
     threading_hook: Callback | None = None
+    use_vi_mode: bool = False
+    vi_mode: ViMode = ViMode.INSERT
 
     ## cached metadata to speed up screen refreshes
     @dataclass
@@ -281,6 +350,11 @@ def __post_init__(self) -> None:
         self.last_refresh_cache.dimensions = (0, 0)
 
     def collect_keymap(self) -> tuple[tuple[KeySpec, CommandName], ...]:
+        if self.use_vi_mode:
+            if self.vi_mode == ViMode.INSERT:
+                return vi_insert_keymap
+            elif self.vi_mode == ViMode.NORMAL:
+                return vi_normal_keymap
         return default_keymap
 
     def calc_screen(self) -> list[str]:
@@ -433,6 +507,57 @@ def eow(self, p: int | None = None) -> int:
             p += 1
         return p
 
+    def vi_eow(self, p: int | None = None) -> int:
+        """Return the 0-based index of the last character of the word
+        following p most immediately (vi 'e' semantics).
+
+        Unlike eow(), this returns the position ON the last word character,
+        not past it. p defaults to self.pos; word boundaries are determined
+        using self.syntax_table."""
+        if p is None:
+            p = self.pos
+        st = self.syntax_table
+        b = self.buffer
+
+        # If we're already at the end of a word, move past it
+        if (p < len(b) and st.get(b[p], SYNTAX_WORD) == SYNTAX_WORD and
+            (p + 1 >= len(b) or st.get(b[p + 1], SYNTAX_WORD) != SYNTAX_WORD)):
+            p += 1
+
+        # Skip non-word characters to find the start of next word
+        while p < len(b) and st.get(b[p], SYNTAX_WORD) != SYNTAX_WORD:
+            p += 1
+
+        # Move to the last character of this word (not past it)
+        while p + 1 < len(b) and st.get(b[p + 1], SYNTAX_WORD) == SYNTAX_WORD:
+            p += 1
+
+        # Clamp to valid buffer range
+        return min(p, len(b) - 1) if b else 0
+
+    def vi_forward_word(self, p: int | None = None) -> int:
+        """Return the 0-based index of the first character of the next word
+        (vi 'w' semantics).
+
+        Unlike eow(), this lands ON the first character of the next word,
+        not past it. p defaults to self.pos; word boundaries are determined
+        using self.syntax_table."""
+        if p is None:
+            p = self.pos
+        st = self.syntax_table
+        b = self.buffer
+
+        # Skip the rest of the current word if we're on one
+        while p < len(b) and st.get(b[p], SYNTAX_WORD) == SYNTAX_WORD:
+            p += 1
+
+        # Skip non-word characters to find the start of next word
+        while p < len(b) and st.get(b[p], SYNTAX_WORD) != SYNTAX_WORD:
+            p += 1
+
+        # Clamp to valid buffer range
+        return min(p, len(b) - 1) if b else 0
+
     def bol(self, p: int | None = None) -> int:
         """Return the 0-based index of the line break preceding p most
         immediately.
@@ -458,6 +583,18 @@ def eol(self, p: int | None = None) -> int:
             p += 1
         return p
 
+    def first_non_whitespace(self, p: int | None = None) -> int:
+        """Return the 0-based index of the first non-whitespace character
+        on the current line.
+
+        p defaults to self.pos."""
+        bol_pos = self.bol(p)
+        eol_pos = self.eol(p)
+        pos = bol_pos
+        while pos < eol_pos and self.buffer[pos].isspace() and self.buffer[pos] != '\n':
+            pos += 1
+        return pos
+
     def max_column(self, y: int) -> int:
         """Return the last x-offset for line y"""
         return self.screeninfo[y][0] + sum(self.screeninfo[y][1])
@@ -589,6 +726,8 @@ def prepare(self) -> None:
             self.pos = 0
             self.dirty = True
             self.last_command = None
+            if self.use_vi_mode:
+                self.enter_insert_mode()
             self.calc_screen()
         except BaseException:
             self.restore()
@@ -760,3 +899,38 @@ def bind(self, spec: KeySpec, command: CommandName) -> None:
     def get_unicode(self) -> str:
         """Return the current buffer as a unicode string."""
         return "".join(self.buffer)
+
+    def enter_insert_mode(self) -> None:
+        if self.vi_mode == ViMode.INSERT:
+            return
+
+        self.vi_mode = ViMode.INSERT
+
+        # Switch translator to insert mode keymap
+        self.keymap = self.collect_keymap()
+        self.input_trans = input.KeymapTranslator(
+            self.keymap, invalid_cls="invalid-key", character_cls="self-insert"
+        )
+
+        self.dirty = True
+
+    def enter_normal_mode(self) -> None:
+        if self.vi_mode == ViMode.NORMAL:
+            return
+
+        self.vi_mode = ViMode.NORMAL
+
+        # Switch translator to normal mode keymap
+        self.keymap = self.collect_keymap()
+        self.input_trans = input.KeymapTranslator(
+            self.keymap, invalid_cls="invalid-key", character_cls="invalid-key"
+        )
+
+        # In vi normal mode, cursor should be ON a character, not after the last one
+        # If we're past the end of line, move back to the last character
+        bol_pos = self.bol()
+        eol_pos = self.eol()
+        if self.pos >= eol_pos and eol_pos > bol_pos:
+            self.pos = eol_pos - 1
+
+        self.dirty = True

Lib/_pyrepl/readline.py

@@ -65,7 +65,6 @@
 
 MoreLinesCallable = Callable[[str], bool]
 
-
 __all__ = [
     "add_history",
     "clear_history",
@@ -344,6 +343,10 @@ def do(self) -> None:
 # ____________________________________________________________
 
 
+def _is_vi_mode_enabled() -> bool:
+    return os.environ.get("PYREPL_VI_MODE", "").lower() in {"1", "true", "on", "yes"}
+
+
 @dataclass(slots=True)
 class _ReadlineWrapper:
     f_in: int = -1
@@ -362,7 +365,11 @@ def __post_init__(self) -> None:
     def get_reader(self) -> ReadlineAlikeReader:
         if self.reader is None:
             console = Console(self.f_in, self.f_out, encoding=ENCODING)
-            self.reader = ReadlineAlikeReader(console=console, config=self.config)
+            self.reader = ReadlineAlikeReader(
+                console=console,
+                config=self.config,
+                use_vi_mode=_is_vi_mode_enabled()
+            )
         return self.reader
 
     def input(self, prompt: object = "") -> str:

Lib/_pyrepl/types.py

@@ -1,3 +1,4 @@
+import enum
 from collections.abc import Callable, Iterator
 
 type Callback = Callable[[], object]
@@ -8,3 +9,8 @@
 type Completer = Callable[[str, int], str | None]
 type CharBuffer = list[str]
 type CharWidths = list[int]
+
+
+class ViMode(str, enum.Enum):
+    INSERT = "insert"
+    NORMAL = "normal"