Implement vi mode and basic motions in pyrepl

Author: lkollar

Lib/_pyrepl/commands.py

@@ -325,20 +325,41 @@ 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.editor_mode.is_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):
     def do(self) -> None:
         self.reader.pos = self.reader.bol()
 
 
+class first_non_whitespace_character(MotionCommand):
+    def do(self) -> None:
+        self.reader.pos = self.reader.first_non_whitespace()
+
+
 class end_of_line(MotionCommand):
     def do(self) -> None:
-        self.reader.pos = self.reader.eol()
+        r = self.reader
+        eol_pos = r.eol()
+        if r.editor_mode.is_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):
@@ -365,6 +386,20 @@ def do(self) -> None:
             r.pos = r.bow()
 
 
+class end_of_word(MotionCommand):
+    def do(self) -> None:
+        r = self.reader
+        for _ in range(r.get_arg()):
+            r.pos = r.vi_eow()
+
+
+class vi_forward_word(MotionCommand):
+    def do(self) -> None:
+        r = self.reader
+        for _ in range(r.get_arg()):
+            r.pos = r.vi_forward_word()
+
+
 class self_insert(EditCommand):
     def do(self) -> None:
         r = self.reader
@@ -503,3 +538,55 @@ def do(self) -> None:
         )
         self.reader.insert(data.replace(done, ""))
         self.reader.last_refresh_cache.invalidated = True
+
+
+class vi_normal_mode(Command):
+    def do(self) -> None:
+        self.reader.enter_normal_mode()
+
+
+class vi_insert_mode(Command):
+    def do(self) -> None:
+        self.reader.enter_insert_mode()
+
+
+class vi_append_mode(Command):
+    def do(self) -> None:
+        if self.reader.pos < len(self.reader.buffer):
+            self.reader.pos += 1
+        self.reader.enter_insert_mode()
+
+
+class vi_append_eol(Command):
+    def do(self) -> None:
+        while self.reader.pos < len(self.reader.buffer):
+            if self.reader.buffer[self.reader.pos] == '\n':
+                break
+            self.reader.pos += 1
+        self.reader.enter_insert_mode()
+
+
+class vi_insert_bol(Command):
+    def do(self) -> None:
+        self.reader.pos = self.reader.first_non_whitespace()
+        self.reader.enter_insert_mode()
+
+
+class vi_open_below(Command):
+    def do(self) -> None:
+        while self.reader.pos < len(self.reader.buffer):
+            if self.reader.buffer[self.reader.pos] == '\n':
+                break
+            self.reader.pos += 1
+
+        self.reader.insert('\n')
+        self.reader.enter_insert_mode()
+
+class vi_open_above(Command):
+    def do(self) -> None:
+        while self.reader.pos > 0 and self.reader.buffer[self.reader.pos - 1] != '\n':
+            self.reader.pos -= 1
+
+        self.reader.insert('\n')
+        self.reader.pos -= 1
+        self.reader.enter_insert_mode()

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()