From 5e45d33606dfcfa32cb3a6bebbf51bc7414ddb8e Mon Sep 17 00:00:00 2001
From: Lie Ryan <lie.1296@gmail.com>
Date: Sun, 1 Mar 2020 22:07:22 +1100
Subject: [PATCH] Set cursor position after replacement to last selected range

---
 pyvim/commands/commands.py | 4 +++-
 tests/test_substitute.py   | 4 ++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/pyvim/commands/commands.py b/pyvim/commands/commands.py
index 7cadb55..826d458 100644
--- a/pyvim/commands/commands.py
+++ b/pyvim/commands/commands.py
@@ -713,6 +713,7 @@ def substitute(editor, range_start, range_end, search, replace, flags):
     """ Substitute /search/ with /replace/ over a range of text """
     def get_line_index_iterator(cursor_position_row, range_start, range_end):
         if not range_start:
+            assert not range_end
             range_start = range_end = cursor_position_row
         else:
             range_start = int(range_start) - 1
@@ -739,7 +740,8 @@ def get_transform_callback(search, replace, flags):
     transform_callback = get_transform_callback(search, replace, flags)
     new_text = buffer.transform_lines(line_index_iterator, transform_callback)
 
-    new_cursor_position_row = int(range_start) - 1 if range_start and not range_end else cursor_position_row
+    assert len(line_index_iterator) >= 1
+    new_cursor_position_row = line_index_iterator[-1]
 
     # update text buffer
     buffer.document = Document(
diff --git a/tests/test_substitute.py b/tests/test_substitute.py
index e0b72c9..5c49e0a 100644
--- a/tests/test_substitute.py
+++ b/tests/test_substitute.py
@@ -55,8 +55,12 @@ def test_substitute_range(editor, editor_buffer):
     assert 'Violet is blue,' in editor_buffer.buffer.text
     assert 'And so are you.' in editor_buffer.buffer.text
     # FIXME: vim would have set the cursor position on last substituted line
+    #        but we set the cursor position on the end_range even when there
+    #        is not substitution there
     # assert editor_buffer.buffer.cursor_position \
     #    == editor_buffer.buffer.text.index('Violet')
+    assert editor_buffer.buffer.cursor_position \
+        == editor_buffer.buffer.text.index('Sugar')
 
 
 def test_substitute_range_boundaries(editor, editor_buffer):