From 70dd8bf24b6577721a19391a40787d7304fed4e1 Mon Sep 17 00:00:00 2001
From: Noorhteen Raja NJ <jnoortheen@gmail.com>
Date: Sat, 8 Jan 2022 04:04:52 +0530
Subject: [PATCH] fix: remove custom handling of osc tokens for ptk prompt
 (#4629)

https://github.com/prompt-toolkit/python-prompt-toolkit/blob/master/examples/prompts/finalterm-shell-integration.py

fixes #374
---
 news/fix-ptk-osc-in-prompt.rst | 24 ++++++++++++++++++++++++
 tests/test_ptk_shell.py        | 29 +----------------------------
 xonsh/ptk_shell/shell.py       | 27 +--------------------------
 3 files changed, 26 insertions(+), 54 deletions(-)
 create mode 100644 news/fix-ptk-osc-in-prompt.rst

diff --git a/news/fix-ptk-osc-in-prompt.rst b/news/fix-ptk-osc-in-prompt.rst
new file mode 100644
index 000000000..c668a298e
--- /dev/null
+++ b/news/fix-ptk-osc-in-prompt.rst
@@ -0,0 +1,24 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* OSC codes in ``$PROMPT`` is no longer removed when using ptk shell.
+  These codes need to be escaped with ``\001..\002`` instead.
+
+**Security:**
+
+* <news item>
diff --git a/tests/test_ptk_shell.py b/tests/test_ptk_shell.py
index b4e075ed1..faac16266 100644
--- a/tests/test_ptk_shell.py
+++ b/tests/test_ptk_shell.py
@@ -7,7 +7,7 @@ import pyte
 
 # verify error if ptk not installed or below min
 
-from xonsh.ptk_shell.shell import tokenize_ansi, remove_ansi_osc
+from xonsh.ptk_shell.shell import tokenize_ansi
 from xonsh.shell import Shell
 
 
@@ -99,33 +99,6 @@ def test_tokenize_ansi(prompt_tokens, ansi_string_parts):
         assert token[1] == text
 
 
-@pytest.mark.parametrize(
-    "raw_prompt, prompt, osc_tokens",
-    [
-        # no title
-        ("test prompt", "test prompt", []),
-        # starts w/ title
-        ("\033]0;TITLE THIS\007test prompt", "test prompt", ["\033]0;TITLE THIS\007"]),
-        # ends w/ title
-        ("test prompt\033]0;TITLE THIS\007", "test prompt", ["\033]0;TITLE THIS\007"]),
-        # title in the middle
-        ("test \033]0;TITLE THIS\007prompt", "test prompt", ["\033]0;TITLE THIS\007"]),
-        # title + iTerm2 OSC exapmle
-        (
-            "test \033]0;TITLE THIS\007prompt \033]133;A\007here",
-            "test prompt here",
-            ["\033]0;TITLE THIS\007", "\033]133;A\007"],
-        ),
-    ],
-)
-def test_remove_ansi_osc(raw_prompt, prompt, osc_tokens):
-    checked_prompt, removed_osc = remove_ansi_osc(raw_prompt)
-    assert prompt == checked_prompt
-    assert len(removed_osc) == len(osc_tokens)
-    for removed, ref in zip(removed_osc, osc_tokens):
-        assert removed == ref
-
-
 @pytest.mark.parametrize(
     "line, exp",
     [
diff --git a/xonsh/ptk_shell/shell.py b/xonsh/ptk_shell/shell.py
index 27ff96598..6cd1aa9ba 100644
--- a/xonsh/ptk_shell/shell.py
+++ b/xonsh/ptk_shell/shell.py
@@ -43,7 +43,6 @@ try:
 except ImportError:
     HAVE_SYS_CLIPBOARD = False
 
-ANSI_OSC_PATTERN = re.compile("\x1b].*?\007")
 CAPITAL_PATTERN = re.compile(r"([a-z])([A-Z])")
 Token = _TokenType()
 
@@ -84,19 +83,6 @@ def tokenize_ansi(tokens):
     return ansi_tokens
 
 
-def remove_ansi_osc(prompt):
-    """Removes the ANSI OSC escape codes - ``prompt_toolkit`` does not support them.
-    Some terminal emulators - like iTerm2 - uses them for various things.
-
-    See: https://www.iterm2.com/documentation-escape-codes.html
-    """
-
-    osc_tokens = ANSI_OSC_PATTERN.findall(prompt)
-    prompt = ANSI_OSC_PATTERN.sub("", prompt)
-
-    return prompt, osc_tokens
-
-
 def _pygments_token_to_classname(token):
     """Converts pygments Tokens, token names (strings) to PTK style names."""
     if token and isinstance(token, str):
@@ -358,7 +344,6 @@ class PromptToolkitShell(BaseShell):
             "refresh_interval": refresh_interval,
             "complete_in_thread": complete_in_thread,
         }
-
         if env["ENABLE_ASYNC_PROMPT"]:
             # once the prompt is done, update it in background as each future is completed
             prompt_args["pre_run"] = self.prompt_formatter.start_update
@@ -427,16 +412,6 @@ class PromptToolkitShell(BaseShell):
         except Exception:  # pylint: disable=broad-except
             print_exception()
 
-        p, osc_tokens = remove_ansi_osc(p)
-
-        if kwargs.get("handle_osc_tokens"):
-            # handle OSC tokens
-            for osc in osc_tokens:
-                if osc[2:4] == "0;":
-                    env["TITLE"] = osc[4:-1]
-                else:
-                    print(osc, file=sys.__stdout__, flush=True)
-
         toks = partial_color_tokenize(p)
 
         return tokenize_ansi(PygmentsTokens(toks))
@@ -447,7 +422,7 @@ class PromptToolkitShell(BaseShell):
             carriage_return()
             self._first_prompt = False
 
-        tokens = self._get_prompt_tokens("PROMPT", "message", handle_osc_tokens=True)
+        tokens = self._get_prompt_tokens("PROMPT", "message")
         self.settitle()
         return tokens