diff --git a/python_files/pythonrc.py b/python_files/pythonrc.py index 182fea98a2c9..e06da5b664ae 100644 --- a/python_files/pythonrc.py +++ b/python_files/pythonrc.py @@ -1,101 +1,101 @@ -import platform -import sys -from enum import Enum - -if sys.platform != "win32": - import readline - -original_ps1 = ">>> " -is_wsl = "microsoft-standard-WSL" in platform.release() - - -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 - - -class REPLHooks: - def __init__(self): - self.global_exit = None - self.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.original_displayhook(value) - - def vscode_excepthook(self, type_, value, traceback): - self.global_exit = value - self.failure_flag = True - self.original_excepthook(type_, value, traceback) +def _initialize(): + import platform + import sys + from enum import Enum - -def get_last_command(): - # Get the last history item - last_command = "" if sys.platform != "win32": - last_command = readline.get_history_item(readline.get_current_history_length()) - return last_command - - -class PS1: - hooks = REPLHooks() - - # 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 - # Guide following official VS Code doc for shell integration sequence: - result = "" - # For non-windows allow recent_command history. + import readline + + original_ps1 = ">>> " + is_wsl = "microsoft-standard-WSL" in platform.release() + + 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 + + class REPLHooks: + def __init__(self): + self.global_exit = None + self.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.original_displayhook(value) + + def vscode_excepthook(self, type_, value, traceback): + self.global_exit = value + self.failure_flag = True + self.original_excepthook(type_, value, traceback) + + def get_last_command(): + # Get the last history item + last_command = "" 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, - prompt=original_ps1, - command_start=ShellIntegrationSequence.COMMAND_START, - ) - 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, - prompt=original_ps1, - command_start=ShellIntegrationSequence.COMMAND_START, - command_executed=ShellIntegrationSequence.COMMAND_EXECUTED, - ) - - return result - - def __repr__(self): - return "" - - -if sys.platform != "win32" and (not is_wsl): - sys.ps1 = PS1() - -ctrl_key = "Cmd" if sys.platform == "darwin" else "Ctrl" - -print(f"{ctrl_key} click to launch VS Code Native REPL (https://aka.ms/python-native-repl)") + last_command = readline.get_history_item(readline.get_current_history_length()) + return last_command + + class PS1: + hooks = REPLHooks() + + # 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 + # Guide following official VS Code doc for shell integration sequence: + result = "" + # For non-windows allow recent_command history. + 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, + prompt=original_ps1, + command_start=ShellIntegrationSequence.COMMAND_START, + ) + 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, + prompt=original_ps1, + command_start=ShellIntegrationSequence.COMMAND_START, + command_executed=ShellIntegrationSequence.COMMAND_EXECUTED, + ) + + return result + + def __repr__(self): + return "" + + if sys.platform != "win32" and (not is_wsl): + sys.ps1 = PS1() + + ctrl_key = "Cmd" if sys.platform == "darwin" else "Ctrl" + + print(f"{ctrl_key} click to launch VS Code Native REPL (https://aka.ms/python-native-repl)") + + +_initialize() +del _initialize diff --git a/python_files/tests/test_shell_integration.py b/python_files/tests/test_shell_integration.py index f1750b3853d4..9902b72f3408 100644 --- a/python_files/tests/test_shell_integration.py +++ b/python_files/tests/test_shell_integration.py @@ -1,6 +1,7 @@ import importlib import platform import sys +from typing import Protocol, cast from unittest.mock import Mock import pythonrc @@ -8,42 +9,43 @@ is_wsl = "microsoft-standard-WSL" in platform.release() +class _Hooks(Protocol): + failure_flag: bool + + +class _PS1(Protocol): + hooks: _Hooks + + def test_decoration_success(): importlib.reload(pythonrc) - ps1 = pythonrc.PS1() - - ps1.hooks.failure_flag = False - result = str(ps1) 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" ) - else: - pass def test_decoration_failure(): importlib.reload(pythonrc) - ps1 = pythonrc.PS1() - - ps1.hooks.failure_flag = True - result = str(ps1) 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" ) - else: - pass def test_displayhook_call(): importlib.reload(pythonrc) - pythonrc.PS1() mock_displayhook = Mock() - hooks = pythonrc.REPLHooks() + hooks = sys.displayhook.__self__ hooks.original_displayhook = mock_displayhook hooks.vscode_displayhook("mock_value") @@ -53,16 +55,21 @@ def test_displayhook_call(): def test_excepthook_call(): importlib.reload(pythonrc) - pythonrc.PS1() mock_excepthook = Mock() - hooks = pythonrc.REPLHooks() + hooks = sys.excepthook.__self__ hooks.original_excepthook = mock_excepthook hooks.vscode_excepthook("mock_type", "mock_value", "mock_traceback") mock_excepthook.assert_called_once_with("mock_type", "mock_value", "mock_traceback") +def test_does_not_pollute_namespace(): + importlib.reload(pythonrc) + + assert not [name for name in vars(pythonrc) if not name.startswith("__")] + + if sys.platform == "darwin": def test_print_statement_darwin(monkeypatch):