ANSI OSC escape code 'support'

xonsh/ptk_shell/shell.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 """The prompt_toolkit based xonsh shell."""
 import os
+import re
 import sys
 import builtins
 from types import MethodType
@@ -37,6 +38,7 @@ from prompt_toolkit.styles.pygments import (
 )
 
 
+ANSI_OSC_PATTERN = re.compile("\x1b].*?\007")
 Token = _TokenType()
 
 events.transmogrify("on_ptk_create", "LoadEvent")
@@ -76,6 +78,25 @@ 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 = re.sub(ANSI_OSC_PATTERN, "", prompt)
+
+    return prompt, osc_tokens
+
+
+def write_ansi_osc(ansi_osc_code):
+    with open(1, "wb", closefd=False) as output:
+        output.write(ansi_osc_code.encode())
+        output.flush()
+
+
 class PromptToolkitShell(BaseShell):
     """The xonsh shell for prompt_toolkit v2 and later."""
 
@@ -250,10 +271,22 @@ class PromptToolkitShell(BaseShell):
             p = self.prompt_formatter(p)
         except Exception:  # pylint: disable=broad-except
             print_exception()
+
+        p, osc_tokens = remove_ansi_osc(p)
+
         toks = partial_color_tokenize(p)
         if self._first_prompt:
             carriage_return()
             self._first_prompt = False
+
+        # handle OSC tokens
+        for osc in osc_tokens:
+            if osc[2:4] == "0;":
+                # set $TITLE
+                builtins.__xonsh__.env["TITLE"] = osc[4:-1]
+            else:
+                write_ansi_osc(osc)
+
         self.settitle()
         return tokenize_ansi(PygmentsTokens(toks))