diff --git a/python_files/pythonrc.py b/python_files/pythonrc.py index 817add379162..1962c92a5d41 100644 --- a/python_files/pythonrc.py +++ b/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,23 +30,41 @@ 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 @@ -47,12 +72,12 @@ def __init__(self): 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 "" - 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" diff --git a/python_files/tests/test_shell_integration.py b/python_files/tests/test_shell_integration.py index 73586601bd01..8d1ef3b4a7f1 100644 --- a/python_files/tests/test_shell_integration.py +++ b/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":