Implement undo

Author: lkollar

Lib/_pyrepl/commands.py

@@ -61,6 +61,8 @@ def do(self) -> None:
 
 
 class KillCommand(Command):
+    modifies_buffer = True
+
     def kill_range(self, start: int, end: int) -> None:
         if start == end:
             return
@@ -422,6 +424,8 @@ def do(self) -> None:
 
 
 class backspace(EditCommand):
+    modifies_buffer = True
+
     def do(self) -> None:
         r = self.reader
         b = r.buffer
@@ -435,6 +439,8 @@ def do(self) -> None:
 
 
 class delete(EditCommand):
+    modifies_buffer = True
+
     def do(self) -> None:
         r = self.reader
         b = r.buffer

Lib/_pyrepl/reader.py

@@ -36,11 +36,12 @@
 
 # types
 Command = commands.Command
-from .types import Callback, SimpleContextManager, KeySpec, CommandName
+from .types import Callback, SimpleContextManager, KeySpec, CommandName, ViUndoState
 
 
 # syntax classes
 SYNTAX_WHITESPACE, SYNTAX_WORD, SYNTAX_SYMBOL = range(3)
+MAX_VI_UNDO_STACK_SIZE = 100
 
 
 from .types import ViMode, ViFindState
@@ -192,6 +193,9 @@ def make_default_commands() -> dict[CommandName, type[Command]]:
         # Replace commands
         (r"r", "vi-replace-char"),
 
+        # Undo commands
+        (r"u", "vi-undo"),
+
         # Special keys still work in normal mode
         (r"\<left>", "left"),
         (r"\<right>", "right"),
@@ -310,6 +314,7 @@ class Reader:
     use_vi_mode: bool = False
     vi_mode: ViMode = ViMode.INSERT
     vi_find: ViFindState = field(default_factory=ViFindState)
+    undo_stack: list[ViUndoState] = field(default_factory=list)
 
     ## cached metadata to speed up screen refreshes
     @dataclass
@@ -894,6 +899,7 @@ def prepare(self) -> None:
             self.last_command = None
             if self.use_vi_mode:
                 self.enter_insert_mode()
+                self.undo_stack.clear()
             self.calc_screen()
         except BaseException:
             self.restore()
@@ -958,6 +964,13 @@ def do_cmd(self, cmd: tuple[str, list[str]]) -> None:
             return  # nothing to do
 
         command = command_type(self, *cmd)  # type: ignore[arg-type]
+
+        # Save undo state in vi mode if the command modifies the buffer
+        if self.use_vi_mode and getattr(command_type, 'modifies_buffer', False):
+            self.undo_stack.append(ViUndoState(
+                buffer_snapshot=self.buffer.copy(),
+                pos_snapshot=self.pos,
+            ))
         command.do()
 
         self.after_command(command)
@@ -1072,6 +1085,14 @@ def enter_insert_mode(self) -> None:
 
         self.vi_mode = ViMode.INSERT
 
+        if len(self.undo_stack) > MAX_VI_UNDO_STACK_SIZE:
+            self.undo_stack.pop(0)
+
+        self.undo_stack.append(ViUndoState(
+            buffer_snapshot=self.buffer.copy(),
+            pos_snapshot=self.pos,
+        ))
+
         # Switch translator to insert mode keymap
         self.keymap = self.collect_keymap()
         self.input_trans = input.KeymapTranslator(

Lib/_pyrepl/types.py

@@ -29,3 +29,8 @@ class ViFindState:
     last_inclusive: bool = True  # f/F=True, t/T=False
     pending_direction: ViFindDirection | None = None
     pending_inclusive: bool = True
+
+@dataclass
+class ViUndoState:
+    buffer_snapshot: CharBuffer = field(default_factory=list)
+    pos_snapshot: int = 0

Lib/_pyrepl/vi_commands.py

@@ -5,6 +5,7 @@
 from .commands import Command, MotionCommand, KillCommand
 from . import input as _input
 from .types import ViFindDirection
+from .trace import trace
 
 
 # ============================================================================
@@ -361,6 +362,8 @@ def do(self) -> None:
 
 class vi_replace_execute(Command):
     """Execute character replacement with the pressed character."""
+    modifies_buffer = True
+
     def do(self) -> None:
         r = self.reader
         r.pop_input_trans()
@@ -375,3 +378,20 @@ class vi_replace_cancel(Command):
     def do(self) -> None:
         r = self.reader
         r.pop_input_trans()
+
+
+# ============================================================================
+# Undo Commands
+# ============================================================================
+
+class vi_undo(Command):
+    """Undo last change (u)."""
+    def do(self) -> None:
+        r = self.reader
+        trace("vi_undo: undo_stack size =", len(r.undo_stack))
+        if not r.undo_stack:
+            return
+        state = r.undo_stack.pop()
+        r.buffer[:] = state.buffer_snapshot
+        r.pos = state.pos_snapshot
+        r.dirty = True

Lib/test/test_pyrepl/test_reader.py

@@ -1173,6 +1173,85 @@ def test_replace_char(self):
         self.assertEqual(reader.pos, 2)  # cursor stays on replaced char
         self.assertEqual(reader.vi_mode, ViMode.NORMAL)  # stays in normal mode
 
+    def test_undo_after_insert(self):
+        events = itertools.chain(
+            code_to_events("hello"),
+            [
+                Event(evt="key", data="\x1b", raw=bytearray(b"\x1b")),  # ESC
+                Event(evt="key", data="u", raw=bytearray(b"u")),        # undo
+            ],
+        )
+        reader, _ = self._run_vi(events)
+        self.assertEqual(reader.get_unicode(), "")
+        self.assertEqual(reader.vi_mode, ViMode.NORMAL)
+
+    def test_undo_after_delete_word(self):
+        events = itertools.chain(
+            code_to_events("hello world"),
+            [
+                Event(evt="key", data="\x1b", raw=bytearray(b"\x1b")),  # ESC
+                Event(evt="key", data="0", raw=bytearray(b"0")),        # BOL
+                Event(evt="key", data="d", raw=bytearray(b"d")),        # delete
+                Event(evt="key", data="w", raw=bytearray(b"w")),        # word
+                Event(evt="key", data="u", raw=bytearray(b"u")),        # undo
+            ],
+        )
+        reader, _ = self._run_vi(events)
+        self.assertEqual(reader.get_unicode(), "hello world")
+
+    def test_undo_after_x_delete(self):
+        events = itertools.chain(
+            code_to_events("hello"),
+            [
+                Event(evt="key", data="\x1b", raw=bytearray(b"\x1b")),  # ESC
+                Event(evt="key", data="x", raw=bytearray(b"x")),        # delete char
+                Event(evt="key", data="u", raw=bytearray(b"u")),        # undo
+            ],
+        )
+        reader, _ = self._run_vi(events)
+        self.assertEqual(reader.get_unicode(), "hello")
+
+    def test_undo_stack_cleared_on_prepare(self):
+        events = itertools.chain(
+            code_to_events("hello"),
+            [Event(evt="key", data="\x1b", raw=bytearray(b"\x1b"))],
+        )
+        reader, console = self._run_vi(events)
+        self.assertGreater(len(reader.undo_stack), 0)
+        reader.prepare()
+        self.assertEqual(len(reader.undo_stack), 0)
+
+    def test_multiple_undo(self):
+        events = itertools.chain(
+            code_to_events("a"),
+            [
+                Event(evt="key", data="\x1b", raw=bytearray(b"\x1b")),  # ESC
+                Event(evt="key", data="a", raw=bytearray(b"a")),        # append
+            ],
+            code_to_events("b"),
+            [
+                Event(evt="key", data="\x1b", raw=bytearray(b"\x1b")),  # ESC
+                Event(evt="key", data="u", raw=bytearray(b"u")),        # undo 'b'
+                Event(evt="key", data="u", raw=bytearray(b"u")),        # undo 'a'
+            ],
+        )
+        reader, _ = self._run_vi(events)
+        self.assertEqual(reader.get_unicode(), "")
+
+    def test_undo_after_replace(self):
+        events = itertools.chain(
+            code_to_events("hello"),
+            [
+                Event(evt="key", data="\x1b", raw=bytearray(b"\x1b")),  # ESC
+                Event(evt="key", data="0", raw=bytearray(b"0")),        # BOL
+                Event(evt="key", data="r", raw=bytearray(b"r")),        # replace
+                Event(evt="key", data="X", raw=bytearray(b"X")),        # replacement
+                Event(evt="key", data="u", raw=bytearray(b"u")),        # undo
+            ],
+        )
+        reader, _ = self._run_vi(events)
+        self.assertEqual(reader.get_unicode(), "hello")
+
 
 @force_not_colorized_test_class
 class TestHistoricalReaderBindings(TestCase):