Restore undo grouping by honoring explicit save_before callbacks

src/prompt_toolkit/key_binding/key_bindings.py

@@ -250,7 +250,7 @@ def add(
         filter: FilterOrBool = True,
         eager: FilterOrBool = False,
         is_global: FilterOrBool = False,
-        save_before: Callable[[KeyPressEvent], bool] = (lambda e: True),
+        save_before: Callable[[KeyPressEvent], bool] | None = None,
         record_in_macro: FilterOrBool = True,
     ) -> Callable[[T], T]:
         """
@@ -267,7 +267,10 @@ def add(
             `Control`, make it a global (always active) binding.
         :param save_before: Callable that takes an `Event` and returns True if
             we should save the current buffer, before handling the event.
-            (That's the default.)
+            When omitted (or `None`), a plain handler saves the buffer before
+            every event, and an existing :class:`Binding` keeps its own
+            `save_before` callable. An explicit callable is used in both
+            cases, also when adding an existing :class:`Binding`.
         :param record_in_macro: Record these key bindings when a macro is
             being recorded. (True by default.)
         """
@@ -294,7 +297,9 @@ def decorator(func: T) -> T:
                             filter=func.filter & to_filter(filter),
                             eager=to_filter(eager) | func.eager,
                             is_global=to_filter(is_global) | func.is_global,
-                            save_before=func.save_before,
+                            save_before=(
+                                func.save_before if save_before is None else save_before
+                            ),
                             record_in_macro=func.record_in_macro,
                         )
                     )
@@ -306,7 +311,9 @@ def decorator(func: T) -> T:
                             filter=filter,
                             eager=eager,
                             is_global=is_global,
-                            save_before=save_before,
+                            save_before=(
+                                (lambda e: True) if save_before is None else save_before
+                            ),
                             record_in_macro=record_in_macro,
                         )
                     )

tests/test_cli.py

@@ -927,6 +927,34 @@ def test_vi_macros():
     assert result.text == "helloworld"
 
 
+def test_vi_undo_removes_whole_insert():
+    feed = partial(_feed_cli_with_input, editing_mode=EditingMode.VI)
+
+    # Insert "asdf", Escape, then undo.
+    result, cli = feed("asdf\x1bu\r")
+    assert result.text == ""
+
+    # An earlier insert stays when the later insert is undone.
+    result, cli = feed("x \x1bAasdf\x1bu\r")
+    assert result.text == "x "
+
+
+def test_emacs_undo_removes_whole_insert():
+    # Insert "asdf", then ControlUnderscore (undo).
+    result, cli = _feed_cli_with_input("asdf\x1f\r")
+    assert result.text == ""
+
+
+def test_emacs_undo_repeated_backspace():
+    # Repeated backspaces are undone as one group.
+    result, cli = _feed_cli_with_input("abc\x7f\x7f\x1f\r")
+    assert result.text == "abc"
+
+    # The second undo removes the insert.
+    result, cli = _feed_cli_with_input("abcd\x7f\x7f\x1f\x1f\r")
+    assert result.text == ""
+
+
 def test_accept_default():
     """
     Test `prompt(accept_default=True)`.

tests/test_key_binding.py

@@ -6,8 +6,9 @@
 
 from prompt_toolkit.application import Application
 from prompt_toolkit.application.current import set_app
+from prompt_toolkit.filters import Condition
 from prompt_toolkit.input.defaults import create_pipe_input
-from prompt_toolkit.key_binding.key_bindings import KeyBindings
+from prompt_toolkit.key_binding.key_bindings import Binding, KeyBindings
 from prompt_toolkit.key_binding.key_processor import KeyPress, KeyProcessor
 from prompt_toolkit.keys import Keys
 from prompt_toolkit.layout import Layout, Window
@@ -202,3 +203,131 @@ def handler(event):
         assert events[1].previous_key_sequence[0].data == "a"
         assert events[1].previous_key_sequence[1].key == "a"
         assert events[1].previous_key_sequence[1].data == "a"
+
+
+def _save_true(event):
+    return True
+
+
+def _save_false(event):
+    return False
+
+
+class _FalsyCallable:
+    """A save_before callable that is itself falsy."""
+
+    def __bool__(self):
+        return False
+
+    def __call__(self, event):
+        return False
+
+
+def _source_binding(save_before, **kwargs):
+    def handler(event):
+        pass
+
+    source_bindings = KeyBindings()
+    source_bindings.add("a", save_before=save_before, **kwargs)(handler)
+    return source_bindings.bindings[0]
+
+
+def _add_binding(binding, **kwargs):
+    kb = KeyBindings()
+    result = kb.add("b", **kwargs)(binding)
+    assert result is binding
+    assert len(kb.bindings) == 1
+    return kb.bindings[0]
+
+
+@pytest.mark.parametrize("source_save_before", [_save_true, _save_false])
+def test_add_binding_inherits_save_before_when_omitted(source_save_before):
+    source = _source_binding(source_save_before)
+
+    assert _add_binding(source).save_before is source_save_before
+    assert _add_binding(source, save_before=None).save_before is source_save_before
+
+
+@pytest.mark.parametrize(
+    "source_save_before, explicit_save_before",
+    [(_save_true, _save_false), (_save_false, _save_true)],
+)
+def test_add_binding_explicit_save_before_overrides(
+    source_save_before, explicit_save_before
+):
+    source = _source_binding(source_save_before)
+
+    added = _add_binding(source, save_before=explicit_save_before)
+
+    assert added.save_before is explicit_save_before
+    assert added.save_before(None) is explicit_save_before(None)
+    # The source binding is not modified.
+    assert source.save_before is source_save_before
+
+
+def test_add_binding_falsy_save_before_callable_overrides():
+    source = _source_binding(_save_true)
+    falsy = _FalsyCallable()
+
+    added = _add_binding(source, save_before=falsy)
+
+    assert added.save_before is falsy
+    assert source.save_before is _save_true
+
+
+def test_add_binding_override_keeps_other_attributes():
+    source_flag = [True]
+    caller_flag = [True]
+    source = _source_binding(
+        _save_true,
+        filter=Condition(lambda: source_flag[0]),
+        record_in_macro=False,
+    )
+
+    added = _add_binding(
+        source,
+        save_before=_save_false,
+        filter=Condition(lambda: caller_flag[0]),
+        eager=True,
+        is_global=True,
+        record_in_macro=True,
+    )
+
+    assert added.handler is source.handler
+    assert added.keys == ("b",)
+    assert added.save_before is _save_false
+    # Filters are still combined with "&".
+    assert added.filter()
+    source_flag[0] = False
+    assert not added.filter()
+    source_flag[0] = True
+    caller_flag[0] = False
+    assert not added.filter()
+    # Eager and is_global are still combined with "|".
+    assert added.eager()
+    assert added.is_global()
+    # record_in_macro is still inherited from the source binding.
+    assert not added.record_in_macro()
+
+
+def test_add_binding_twice_keeps_explicit_save_before():
+    source = _source_binding(_save_true)
+    first = _add_binding(source, save_before=_save_false)
+
+    assert _add_binding(first).save_before is _save_false
+
+
+def test_add_handler_save_before_defaults():
+    def handler(event):
+        pass
+
+    kb = KeyBindings()
+    assert kb.add("a")(handler) is handler
+    kb.add("b", save_before=None)(handler)
+    kb.add("c", save_before=_save_false)(handler)
+    omitted, explicit_none, explicit = kb.bindings
+
+    assert isinstance(omitted, Binding)
+    assert omitted.save_before(None) is True
+    assert explicit_none.save_before(None) is True
+    assert explicit.save_before is _save_false