From e7426ce70986bae22f60e48db086e9c45978897f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C4=90=E1=BB=97=20Trung=20Nguy=C3=AAn?=
 <dtngn@users.noreply.github.com>
Date: Mon, 27 Jan 2025 15:56:19 +0700
Subject: [PATCH] Add configuring cursor shape for prompt toolkit (#5785)

* env: add `PROMPT_TOOLKIT_CURSOR_SHAPE` for configuring `prompt_toolkit` cursor shape

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* accepting prompt_toolkit.cursor_shapes.CursorShape as value

* rename PROMPT_TOOLKIT_CURSOR_SHAPE => XONSH_PROMPT_CURSOR_SHAPE

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 news/config-cursor-shape.rst       | 23 +++++++++++++
 xonsh/environ.py                   | 15 ++++++++
 xonsh/shells/ptk_shell/__init__.py | 13 +++----
 xonsh/tools.py                     | 55 ++++++++++++++++++++++++++++++
 4 files changed, 97 insertions(+), 9 deletions(-)
 create mode 100644 news/config-cursor-shape.rst

diff --git a/news/config-cursor-shape.rst b/news/config-cursor-shape.rst
new file mode 100644
index 000000000..c183577b8
--- /dev/null
+++ b/news/config-cursor-shape.rst
@@ -0,0 +1,23 @@
+**Added:**
+
+* env: add ``$XONSH_PROMPT_CURSOR_SHAPE`` for configuring prompt cursor shape.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>
diff --git a/xonsh/environ.py b/xonsh/environ.py
index b49cd8f56..1c533f86c 100644
--- a/xonsh/environ.py
+++ b/xonsh/environ.py
@@ -98,6 +98,8 @@ from xonsh.tools import (
     to_int_or_none,
     to_itself,
     to_logfile_opt,
+    to_ptk_cursor_shape,
+    to_ptk_cursor_shape_display_value,
     to_repr_pretty_,
     to_shlvl,
     to_tok_color_dict,
@@ -1743,6 +1745,19 @@ class PTKSetting(PromptSetting):  # sub-classing -> sub-group
         "``DEPTH_1_BIT``, ``DEPTH_4_BIT``, ``DEPTH_8_BIT``, ``DEPTH_24_BIT`` "
         "colors. Default is an empty string which means that prompt toolkit decide.",
     )
+    XONSH_PROMPT_CURSOR_SHAPE = Var(
+        always_false,
+        to_ptk_cursor_shape,
+        to_ptk_cursor_shape_display_value,
+        to_ptk_cursor_shape("modal-vi-mode-only"),
+        "The cursor shape. Possible values for prompt toolkit are: "
+        "``block``, ``beam``, ``underline``, "
+        "``blinking-block``, ``blinking-beam``, ``blinking-underline``, "
+        "``modal``, ``modal-vi-mode-only``, ``never-change``. "
+        "Default value is ``modal-vi-mode-only`` which means "
+        "``modal`` if in vi mode and ``never-change`` if not in vi mode.",
+        doc_default="modal-vi-mode-only",
+    )
     PTK_STYLE_OVERRIDES = Var(
         is_tok_color_dict,
         to_tok_color_dict,
diff --git a/xonsh/shells/ptk_shell/__init__.py b/xonsh/shells/ptk_shell/__init__.py
index c2827e91f..c42429e2b 100644
--- a/xonsh/shells/ptk_shell/__init__.py
+++ b/xonsh/shells/ptk_shell/__init__.py
@@ -46,13 +46,6 @@ try:
 except ImportError:
     HAVE_SYS_CLIPBOARD = False
 
-try:
-    from prompt_toolkit.cursor_shapes import ModalCursorShapeConfig
-
-    HAVE_CURSOR_SHAPE = True
-except ImportError:
-    HAVE_CURSOR_SHAPE = False
-
 CAPITAL_PATTERN = re.compile(r"([a-z])([A-Z])")
 Token = _TokenType()
 
@@ -377,8 +370,10 @@ class PromptToolkitShell(BaseShell):
             for attr, val in self.get_lazy_ptk_kwargs():
                 prompt_args[attr] = val
 
-        if editing_mode == EditingMode.VI and HAVE_CURSOR_SHAPE:
-            prompt_args["cursor"] = ModalCursorShapeConfig()
+        cursor_shape = env.get("XONSH_PROMPT_CURSOR_SHAPE")
+        if cursor_shape:
+            prompt_args["cursor"] = cursor_shape
+
         events.on_pre_prompt.fire()
         line = self.prompter.prompt(**prompt_args)
         events.on_post_prompt.fire()
diff --git a/xonsh/tools.py b/xonsh/tools.py
index 3da5bd8b9..9d496f7f7 100644
--- a/xonsh/tools.py
+++ b/xonsh/tools.py
@@ -41,6 +41,19 @@ import typing as tp
 import warnings
 from contextlib import contextmanager
 
+try:
+    from prompt_toolkit.cursor_shapes import (
+        CursorShape,
+        CursorShapeConfig,
+        DynamicCursorShapeConfig,
+        ModalCursorShapeConfig,
+        SimpleCursorShapeConfig,
+    )
+
+    HAVE_CURSOR_SHAPE = True
+except ImportError:
+    HAVE_CURSOR_SHAPE = False
+
 # adding imports from further xonsh modules is discouraged to avoid circular
 # dependencies
 from xonsh import __version__
@@ -1728,6 +1741,48 @@ def ptk2_color_depth_setter(x):
     return x
 
 
+def ptk_cursor_shape_vi_modal():
+    if xsh.env.get("VI_MODE"):
+        return ModalCursorShapeConfig()
+    else:
+        return SimpleCursorShapeConfig()
+
+
+def to_ptk_cursor_shape(x):
+    if not HAVE_CURSOR_SHAPE:
+        return None
+    if isinstance(x, (CursorShape, CursorShapeConfig)):
+        return x
+    if not isinstance(x, str):
+        raise ValueError("invalid cursor shape")
+    x = str(x).upper().replace("-", "_")
+    if x == "MODAL":
+        return ModalCursorShapeConfig()
+    elif x == "MODAL_VI_MODE_ONLY":
+        return DynamicCursorShapeConfig(ptk_cursor_shape_vi_modal)
+    try:
+        return CursorShape[x]
+    except KeyError:
+        return SimpleCursorShapeConfig()
+
+
+def to_ptk_cursor_shape_display_value(x):
+    if not x:
+        return ""
+    if isinstance(x, SimpleCursorShapeConfig):
+        x = x.get_cursor_shape(None)
+    if isinstance(x, CursorShape):
+        x = x.value.lower().replace("_", "-")
+        if x.startswith("-"):
+            x = x[1:]
+        return x
+    if isinstance(x, ModalCursorShapeConfig):
+        return "modal"
+    if isinstance(x, DynamicCursorShapeConfig):
+        return "modal-vi-mode-only"
+    return "unknown"
+
+
 def is_completions_display_value(x):
     """Enumerated values of ``$COMPLETIONS_DISPLAY``"""
     return x in {"none", "single", "multi"}