Set cursor x after vertical movement

editor/main.py

@@ -53,22 +53,41 @@ class Buffer:
         self.cx = cx
         self.cy = cy
 
+        self._cx_hint = cx
+
     def up(self) -> "Buffer":
         if self.cy > 0:
             self.cy -= 1
+            self._set_cx_after_vertical_movement()
         return self
 
     def down(self) -> "Buffer":
         if self.cy < len(self._lines) - 1:
             self.cy += 1
+            self._set_cx_after_vertical_movement()
         return self
 
     def left(self) -> "Buffer":
         if self.cx > 0:
             self.cx -= 1
+            self._cx_hint = self.cx
         return self
 
     def right(self) -> "Buffer":
-        if self.cx < len(self._lines[self.cy]) - 1:
+        if self.cx < self._max_cx:
             self.cx += 1
+            self._cx_hint = self.cx
         return self
+
+    def _set_cx_after_vertical_movement(self) -> None:
+        if self.cx > self._max_cx:
+            # Cursor exceeded the line
+            if self.cx > self._cx_hint:
+                self._cx_hint = self.cx
+            self.cx = max(self._max_cx, 0)
+        else:
+            self.cx = min(self._cx_hint, self._max_cx)
+
+    @property
+    def _max_cx(self) -> int:
+        return len(self._lines[self.cy]) - 1

tests/main_test.py

@@ -2,17 +2,99 @@ from editor.main import Buffer
 
 
 def test_buffer_up():
-    assert Buffer(cy=1).up().cy == 0
+    assert Buffer(["foo", "bar"], cy=1).up().cy == 0
 
 
 def test_buffer_up_at_first_line():
-    assert Buffer().up().cy == 0
+    assert Buffer(["foo"]).up().cy == 0
+
+
+def test_buffer_up_passed_shorter_line():
+    # It correctly resets to a shorter and then longer line
+    buf = Buffer(["longer line", "short", "long line"], cx=8, cy=2)
+    buf.up()
+    assert buf.cx == 4
+    assert buf.cy == 1
+    buf.up()
+    assert buf.cx == 8
+    assert buf.cy == 0
+
+    # It correctly resets to a shorter and then medium line
+    buf = Buffer(["long line", "short", "longer line"], cx=10, cy=2)
+    buf.up()
+    assert buf.cx == 4
+    assert buf.cy == 1
+    buf.up()
+    assert buf.cx == 8
+    assert buf.cy == 0
+
+    # It correctly resets cx on empty lines
+    buf = Buffer(["", "foo"], cx=4, cy=1)
+    buf.up()
+    assert buf.cx == 0
+    assert buf.cy == 0
+
+    # It correctly resets cx on lines of the same length
+    buf = Buffer(["short", "short"], cx=4, cy=1)
+    buf.up()
+    assert buf.cx == 4
+    assert buf.cy == 0
+
+    # It correctly resets cx hint after horizontal movement
+    buf = Buffer(["foo", "", "bar"], cy=2)
+    buf.right().up().up()
+    assert buf.cx == 1
+    assert buf.cy == 0
+    buf.left().down().down()
+    assert buf.cx == 0
+    assert buf.cy == 2
 
 
 def test_buffer_down():
     assert Buffer(["foo", "bar"]).down().cy == 1
 
 
+def test_buffer_down_passed_shorter_line():
+    # It correctly resets to a shorter and then longer line
+    buf = Buffer(["long line", "short", "longer line"], cx=8)
+    buf.down()
+    assert buf.cx == 4
+    assert buf.cy == 1
+    buf.down()
+    assert buf.cx == 8
+    assert buf.cy == 2
+
+    # It correctly resets to a shorter and then medium line
+    buf = Buffer(["longer line", "short", "long line"], cx=10)
+    buf.down()
+    assert buf.cx == 4
+    assert buf.cy == 1
+    buf.down()
+    assert buf.cx == 8
+    assert buf.cy == 2
+
+    # It correctly resets cx on empty lines
+    buf = Buffer(["foo", ""], cx=4)
+    buf.down()
+    assert buf.cx == 0
+    assert buf.cy == 1
+
+    # It correctly resets cx on lines of the same length
+    buf = Buffer(["short", "short"], cx=4)
+    buf.down()
+    assert buf.cx == 4
+    assert buf.cy == 1
+
+    # It correctly resets cx hint after horizontal movement
+    buf = Buffer(["foo", "", "bar"])
+    buf.right().down().down()
+    assert buf.cx == 1
+    assert buf.cy == 2
+    buf.left().up().up()
+    assert buf.cx == 0
+    assert buf.cy == 0
+
+
 def test_buffer_down_at_last_line():
     assert Buffer(["foo"]).down().cy == 0