PyREPL: Correctly handle OSC codes

python_files/pythonrc.py

@@ -1,14 +1,21 @@
 def _initialize():
     import platform
     import sys
-    from enum import Enum
 
     if sys.platform != "win32":
         import readline
 
-    original_ps1 = ">>> "
-    is_wsl = "microsoft-standard-WSL" in platform.release()
+    # Override regex to support OSC codes
+    if sys.version_info >= (3, 13):
+        import re
+
+        import _pyrepl.utils
+
+        _pyrepl.utils.ANSI_ESCAPE_SEQUENCE = re.compile(
+            r"\x1b(?:\[[ -@]*[A-~]|\][^\x1b\x07]*(?:\x07|\x1b\\))"
+        )
 
+    original_ps1 = ">>> "
     # PYTHONSTARTUP executes this file's code inside the user's __main__
     # namespace, so PS1.__str__'s globals are the user's globals. If the
     # user later shadows a name we rely on at prompt-render time (e.g.
@@ -23,36 +30,54 @@ def _initialize():
     _bool = bool
     _str = str
 
-    class ShellIntegrationSequence(str, Enum):
-        SOH = "\001"
-        STX = "\002"
-        COMMAND_EXECUTED = "\x1b]633;C\x07"
-        COMMAND_LINE = "\x1b]633;E;"
-        COMMAND_FINISHED = "\x1b]633;D;"
-        PROMPT_STARTED = "\x1b]633;A\x07"
-        COMMAND_START = "\x1b]633;B\x07"
-        TERMINATOR = "\x07"
-
-        def __str__(self):
-            return self.value
+    # https://code.visualstudio.com/docs/terminal/shell-integration#_supported-escape-sequences
+    class ShellIntegrationSequence:
+        soh = "\001"
+        stx = "\002"
+        template = "\x1b]633;{}\x07"
+
+        # Before the prompt (>>>) is displayed
+        @classmethod
+        def prompt_start(cls) -> str:
+            return cls.template.format("A")
+
+        # After the prompt (>>>) is displayed
+        @classmethod
+        def prompt_end(cls) -> str:
+            return cls.template.format("B")
+
+        # After the user has typed a command but before it is executed
+        @classmethod
+        def pre_execution(cls) -> str:
+            return cls.template.format("C")
+
+        @classmethod
+        def execution_finished(cls, exit_code: int) -> str:
+            """Mark execution as finished with its exit code."""
+            return cls.template.format(f"D;{exit_code}")
+
+        @classmethod
+        def command_line(cls, command: object) -> str:
+            """Explicitly set the command line interpreted by the shell."""
+            return cls.template.format(f"E;{command}")
 
     class REPLHooks:
         def __init__(self):
             self.global_exit = None
-            self.failure_flag = False
+            self.last_failure_flag = False
             self.original_excepthook = sys.excepthook
             self.original_displayhook = sys.displayhook
             sys.excepthook = self.vscode_excepthook
             sys.displayhook = self.vscode_displayhook
 
         def vscode_displayhook(self, value):
             if value is None:
-                self.failure_flag = False
+                self.last_failure_flag = False
             self.original_displayhook(value)
 
         def vscode_excepthook(self, type_, value, traceback):
             self.global_exit = value
-            self.failure_flag = True
+            self.last_failure_flag = True
             self.original_excepthook(type_, value, traceback)
 
     def get_last_command():
@@ -67,43 +92,40 @@ class PS1:
 
         # str will get called for every prompt with exit code to show success/failure
         def __str__(self):
-            exit_code = _int(_bool(self.hooks.failure_flag))
-            self.hooks.failure_flag = False
+            last_exit_code = _int(_bool(self.hooks.last_failure_flag))
+            self.hooks.last_failure_flag = False
             # Guide following official VS Code doc for shell integration sequence:
             result = ""
             # For non-windows allow recent_command history.
+            # fmt: off
             if sys.platform != "win32":
-                result = "{soh}{command_executed}{command_line}{command_finished}{prompt_started}{stx}{prompt}{soh}{command_start}{stx}".format(
-                    soh=ShellIntegrationSequence.SOH,
-                    stx=ShellIntegrationSequence.STX,
-                    command_executed=ShellIntegrationSequence.COMMAND_EXECUTED,
-                    command_line=ShellIntegrationSequence.COMMAND_LINE
-                    + _str(get_last_command())
-                    + ShellIntegrationSequence.TERMINATOR,
-                    command_finished=ShellIntegrationSequence.COMMAND_FINISHED
-                    + _str(exit_code)
-                    + ShellIntegrationSequence.TERMINATOR,
-                    prompt_started=ShellIntegrationSequence.PROMPT_STARTED,
+                result = "{soh}{pre_execution}{command_line}{execution_finished}{prompt_start}{stx}{prompt}{soh}{prompt_end}{stx}".format(  # noqa: UP032
+                    soh=ShellIntegrationSequence.soh,
+                    pre_execution=ShellIntegrationSequence.pre_execution(),
+                    command_line=ShellIntegrationSequence.command_line(get_last_command()),
+                    execution_finished=ShellIntegrationSequence.execution_finished(last_exit_code),
+                    prompt_start=ShellIntegrationSequence.prompt_start(),
+                    stx=ShellIntegrationSequence.stx,
                     prompt=original_ps1,
-                    command_start=ShellIntegrationSequence.COMMAND_START,
+                    prompt_end=ShellIntegrationSequence.prompt_end(),
                 )
             else:
-                result = "{command_finished}{prompt_started}{prompt}{command_start}{command_executed}".format(
-                    command_finished=ShellIntegrationSequence.COMMAND_FINISHED
-                    + _str(exit_code)
-                    + ShellIntegrationSequence.TERMINATOR,
-                    prompt_started=ShellIntegrationSequence.PROMPT_STARTED,
+                result = "{execution_finished}{prompt_start}{prompt}{prompt_end}{pre_execution}".format(  # noqa: UP032
+                    execution_finished=ShellIntegrationSequence.execution_finished(last_exit_code),
+                    prompt_start=ShellIntegrationSequence.prompt_start(),
                     prompt=original_ps1,
-                    command_start=ShellIntegrationSequence.COMMAND_START,
-                    command_executed=ShellIntegrationSequence.COMMAND_EXECUTED,
+                    prompt_end=ShellIntegrationSequence.prompt_end(),
+                    pre_execution=ShellIntegrationSequence.pre_execution(),
                 )
+            # fmt: on
 
             return result
 
         def __repr__(self):
             return "<Custom PS1 for VS Code Python Shell Integration>"
 
-    if sys.platform != "win32" and (not is_wsl):
+    is_wsl = "microsoft-standard-WSL" in platform.release()
+    if not is_wsl and (sys.platform != "win32" or sys.version_info >= (3, 13)):
         sys.ps1 = PS1()
 
     ctrl_key = "Cmd" if sys.platform == "darwin" else "Ctrl"

python_files/tests/test_shell_integration.py

@@ -8,40 +8,48 @@
 import pythonrc
 
 is_wsl = "microsoft-standard-WSL" in platform.release()
+prompt_is_installed = not is_wsl and (sys.platform != "win32" or sys.version_info >= (3, 13))
 
 PYTHONRC_PATH = Path(pythonrc.__file__)
 
 
 class _Hooks(Protocol):
-    failure_flag: bool
+    last_failure_flag: bool
 
 
 class _PS1(Protocol):
     hooks: _Hooks
 
 
+def _expected_prompt(exit_code: int) -> str:
+    if sys.platform == "win32":
+        return f"\x1b]633;D;{exit_code}\x07\x1b]633;A\x07>>> \x1b]633;B\x07\x1b]633;C\x07"
+    return (
+        "\x01\x1b]633;C\x07\x1b]633;E;None\x07"
+        f"\x1b]633;D;{exit_code}\x07\x1b]633;A\x07\x02>>> \x01\x1b]633;B\x07\x02"
+    )
+
+
 def test_decoration_success():
     importlib.reload(pythonrc)
-    if sys.platform != "win32" and (not is_wsl):
-        ps1 = cast("_PS1", sys.ps1)
-        ps1.hooks.failure_flag = False
-        result = str(ps1)
-        assert (
-            result
-            == "\x01\x1b]633;C\x07\x1b]633;E;None\x07\x1b]633;D;0\x07\x1b]633;A\x07\x02>>> \x01\x1b]633;B\x07\x02"
-        )
+    if not prompt_is_installed:
+        return
+
+    ps1 = cast("_PS1", sys.ps1)
+    ps1.hooks.last_failure_flag = False
+
+    assert str(ps1) == _expected_prompt(0)
 
 
 def test_decoration_failure():
     importlib.reload(pythonrc)
-    if sys.platform != "win32" and (not is_wsl):
-        ps1 = cast("_PS1", sys.ps1)
-        ps1.hooks.failure_flag = True
-        result = str(ps1)
-        assert (
-            result
-            == "\x01\x1b]633;C\x07\x1b]633;E;None\x07\x1b]633;D;1\x07\x1b]633;A\x07\x02>>> \x01\x1b]633;B\x07\x02"
-        )
+    if not prompt_is_installed:
+        return
+
+    ps1 = cast("_PS1", sys.ps1)
+    ps1.hooks.last_failure_flag = True
+
+    assert str(ps1) == _expected_prompt(1)
 
 
 def test_displayhook_call():
@@ -73,6 +81,40 @@ def test_does_not_pollute_namespace():
     assert not [name for name in vars(pythonrc) if not name.startswith("__")]
 
 
+def test_replacement_regex_removes_bel_terminated_osc():
+    if sys.version_info < (3, 13):
+        return
+
+    pyrepl_utils = importlib.import_module("_pyrepl.utils")
+    importlib.reload(pythonrc)
+
+    assert pyrepl_utils.ANSI_ESCAPE_SEQUENCE.sub("", "before\x1b]633;A\x07after") == ("beforeafter")
+
+
+def test_replacement_regex_removes_st_terminated_osc():
+    if sys.version_info < (3, 13):
+        return
+
+    pyrepl_utils = importlib.import_module("_pyrepl.utils")
+    importlib.reload(pythonrc)
+
+    assert pyrepl_utils.ANSI_ESCAPE_SEQUENCE.sub("", "before\x1b]633;A\x1b\\after") == (
+        "beforeafter"
+    )
+
+
+def test_replacement_regex_preserves_csi_handling():
+    if sys.version_info < (3, 13):
+        return
+
+    pyrepl_utils = importlib.import_module("_pyrepl.utils")
+    importlib.reload(pythonrc)
+
+    assert pyrepl_utils.ANSI_ESCAPE_SEQUENCE.sub("", "before\x1b[31mred\x1b[0mafter") == (
+        "beforeredafter"
+    )
+
+
 def test_prompt_survives_shadowed_builtins_under_pythonstartup():
     # PYTHONSTARTUP executes pythonrc's source directly inside the real
     # REPL's __main__ namespace, not as an imported module. The tests
@@ -81,7 +123,7 @@ def test_prompt_survives_shadowed_builtins_under_pythonstartup():
     # the real PYTHONSTARTUP path by exec-ing the source into a synthetic
     # __main__-like namespace, then shadow the names PS1 relies on at
     # prompt-render time and confirm rendering the prompt still works.
-    if sys.platform == "win32" or is_wsl:
+    if not prompt_is_installed:
         return
 
     source = PYTHONRC_PATH.read_text(encoding="utf-8")
@@ -100,11 +142,7 @@ def test_prompt_survives_shadowed_builtins_under_pythonstartup():
     )
 
     ps1 = cast("_PS1", sys.ps1)
-    result = str(ps1)
-    assert (
-        result
-        == "\x01\x1b]633;C\x07\x1b]633;E;None\x07\x1b]633;D;0\x07\x1b]633;A\x07\x02>>> \x01\x1b]633;B\x07\x02"
-    )
+    assert str(ps1) == _expected_prompt(0)
 
 
 if sys.platform == "darwin":