cursor shape changes with vi editing mode

Author: marskar

ptpython/key_bindings.py

@@ -1,3 +1,5 @@
+import sys
+
 from prompt_toolkit.application import get_app
 from prompt_toolkit.document import Document
 from prompt_toolkit.enums import DEFAULT_BUFFER
@@ -11,6 +13,7 @@
 )
 from prompt_toolkit.key_binding import KeyBindings
 from prompt_toolkit.keys import Keys
+from prompt_toolkit.key_binding.vi_state import InputMode, ViState
 
 from .utils import document_is_multiline_python
 
@@ -319,3 +322,32 @@ def auto_newline(buffer):
         if current_line[-1:] == ":":
             for x in range(4):
                 insert_text(" ")
+
+
+def get_input_mode(self):
+    if sys.version_info[0] == 3:
+        # Decrease input flush timeout from 500ms to 10ms.
+        app = get_app()
+        app.ttimeoutlen = 0.01
+        # Decrease handler call timeout from 1s to 400ms
+        app.timeoutlen = 0.4
+
+    return self._input_mode
+
+
+def set_input_mode(self, mode):
+    shape = {InputMode.NAVIGATION: 2, InputMode.REPLACE: 4}.get(mode, 6)
+    cursor = "\x1b[{} q".format(shape)
+
+    if hasattr(sys.stdout, "_cli"):
+        write = sys.stdout._cli.output.write_raw
+    else:
+        write = sys.stdout.write
+
+    write(cursor)
+    sys.stdout.flush()
+
+    self._input_mode = mode
+
+ViState._input_mode = InputMode.INSERT
+ViState.input_mode = property(get_input_mode, set_input_mode)