From ea4fe8222fef99b562f72818a21877d01d145788 Mon Sep 17 00:00:00 2001 From: le_woudar Date: Mon, 25 Jul 2022 19:42:43 +0200 Subject: [PATCH 1/2] docs: updated README.md and docs/installation.md added reference to prompt_toolkit in the installation section --- README.md | 1 + docs/installation.md | 1 + 2 files changed, 2 insertions(+) diff --git a/README.md b/README.md index 61503e1..8ddc523 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,7 @@ ws starts working from **python3.7** and also supports **pypy3**. It has the fol - [click-didyoumean](https://pypi.org/project/click-didyoumean/) for command suggestions in case of typos. - [rich](https://rich.readthedocs.io/en/latest/) for beautiful output display. - [shellingham](https://pypi.org/project/shellingham/) to detect the shell used. +- [prompt_toolkit](https://python-prompt-toolkit.readthedocs.io/en/master/) for interactive command line. ## Usage diff --git a/docs/installation.md b/docs/installation.md index 22528b9..b3cc65c 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -24,6 +24,7 @@ ws starts working from **python3.7** and also supports **pypy3**. It has the fol - [click-didyoumean](https://pypi.org/project/click-didyoumean/) for command suggestions in case of typos. - [rich](https://rich.readthedocs.io/en/latest/) for beautiful output display. - [shellingham](https://pypi.org/project/shellingham/) to detect the shell used. +- [prompt_toolkit](https://python-prompt-toolkit.readthedocs.io/en/master/) for interactive command line. To confirm the installation works, you can type the command name in the terminal. From 85b2a697e75edc9407febf9455692143e6b56edb Mon Sep 17 00:00:00 2001 From: le_woudar Date: Mon, 25 Jul 2022 23:25:08 +0200 Subject: [PATCH 2/2] feat(completion): implemented powershell cli completion --- tests/{ => commands}/test_completion.py | 72 ++++++++++++++++++++-- ws/commands/completion.py | 79 ++++++++++++++++++++++++- 2 files changed, 146 insertions(+), 5 deletions(-) rename tests/{ => commands}/test_completion.py (59%) diff --git a/tests/test_completion.py b/tests/commands/test_completion.py similarity index 59% rename from tests/test_completion.py rename to tests/commands/test_completion.py index 5f0683a..eec7bc2 100644 --- a/tests/test_completion.py +++ b/tests/commands/test_completion.py @@ -5,10 +5,49 @@ import pytest import shellingham -from ws.commands.completion import SHELLS +from ws.commands.completion import SHELLS, install_powershell from ws.main import cli +class TestInstallPowershell: + """Tests function install_powershell""" + + @pytest.mark.parametrize('shell', ['powershell', 'pwsh']) + def test_should_print_error_when_unable_to_get_path_to_completion_script(self, capsys, mocker, shell): + def fake_subprocess_run(command_args, *_, **_k): + if command_args[1] != '-NoProfile': + return + raise subprocess.CalledProcessError(returncode=1, cmd='pwsh') + + with pytest.raises(SystemExit): + mocker.patch('subprocess.run', side_effect=fake_subprocess_run) + install_powershell(shell) # type: ignore + + @pytest.mark.parametrize('shell', ['powershell', 'pwsh']) + def test_should_create_or_update_user_profile(self, tmp_path, mocker, shell): + user_profile_path = tmp_path / 'WindowsPowerShell' / 'Microsoft.PowerShell_profile.ps1' + + def fake_subprocess_run(command_args, *_, **_k): + if command_args[1] == '-NoProfile': + return subprocess.CompletedProcess(command_args, 0, bytes(user_profile_path), None) + + mocker.patch( + 'shellingham.detect_shell', + return_value=(shell, f'C:\\Windows\\System32\\WindowsPowershell\\v1.0\\{shell}.exe'), + ) + mocker.patch('subprocess.run', side_effect=fake_subprocess_run) + install_powershell(shell) # type: ignore + + # check user profile file + assert user_profile_path.is_file() + + content = user_profile_path.read_text() + assert content.startswith('Import-Module PSReadLine') + assert '$Env:_WS_COMPLETE = "complete_powershell"' in content + assert 'ws | ForEach-Object {' in content + assert content.endswith('Register-ArgumentCompleter -Native -CommandName ws -ScriptBlock $scriptblock\n') + + def test_should_print_error_when_shell_is_not_detected(mocker, runner): mocker.patch('shellingham.detect_shell', side_effect=shellingham.ShellDetectionFailure) result = runner.invoke(cli, ['install-completion']) @@ -27,12 +66,13 @@ def test_should_print_error_when_os_name_is_unknown(monkeypatch, runner): def test_should_print_error_if_shell_is_not_supported(mocker, runner): - mocker.patch('shellingham.detect_shell', return_value=('pwsh', 'C:\\bin\\pwsh')) + mocker.patch('shellingham.detect_shell', return_value=('cmd', 'C:\\bin\\cmd.exe')) result = runner.invoke(cli, ['install-completion']) assert result.exit_code == 1 - shells_string = ', '.join(SHELLS) - assert f'Your shell is not supported. Shells supported are: {shells_string}\n' == result.output + shells_string = ', '.join(SHELLS[:-1]) + assert f'Your shell is not supported. Shells supported are: {shells_string}' in result.output + assert result.output.endswith('pwsh\n') @pytest.mark.parametrize('shell', [('bash', '/bin/bash'), ('zsh', '/bin/zsh'), ('fish', '/bin/fish')]) @@ -116,3 +156,27 @@ def test_should_create_completion_file_and_install_it_for_fish_shell(tmp_path, m content = completion_file.read_text() assert content.startswith('function _ws_completion') assert content.endswith('"(_ws_completion)";\n\n') + + +@pytest.mark.skipif(platform.system() in ['Darwin', 'Linux'], reason='powershell is not supported on these OS') +@pytest.mark.parametrize('shell', ['powershell', 'pwsh']) +def test_should_create_completion_script_and_add_it_in_powershell_profile(tmp_path, mocker, runner, shell): + user_profile_path = tmp_path / 'WindowsPowerShell' / 'Microsoft.PowerShell_profile.ps1' + + def fake_subprocess_run(command_args, *_, **_k): + if command_args[1] == '-NoProfile': + return subprocess.CompletedProcess(command_args, 0, bytes(user_profile_path), None) + + mocker.patch('subprocess.run', side_effect=fake_subprocess_run) + + result = runner.invoke(cli, ['install-completion']) + + assert result.exit_code == 0 + assert 'Successfully installed completion script!\n' in result.output + + # check user profile file + assert user_profile_path.is_file() + + content = user_profile_path.read_text() + assert content.startswith('Import-Module PSReadLine') + assert content.endswith('Register-ArgumentCompleter -Native -CommandName ws -ScriptBlock $scriptblock\n') diff --git a/ws/commands/completion.py b/ws/commands/completion.py index dbac834..1eb1d81 100644 --- a/ws/commands/completion.py +++ b/ws/commands/completion.py @@ -1,12 +1,87 @@ +from __future__ import annotations + +import os import subprocess # nosec from pathlib import Path import click import shellingham +from click.shell_completion import ShellComplete, add_completion_class +from typing_extensions import Literal from ws.console import console -SHELLS = ['bash', 'zsh', 'fish'] +SHELLS = ['bash', 'zsh', 'fish', 'powershell', 'pwsh'] + +# Windows support code is heavily inspired by the typer project +POWERSHELL_COMPLETION_SCRIPT = """ +Import-Module PSReadLine +Set-PSReadLineKeyHandler -Chord Tab -Function MenuComplete +$scriptblock = { + param($wordToComplete, $commandAst, $cursorPosition) + $Env:%(complete_var)s = "complete_powershell" + $Env:_CLICK_COMPLETE_ARGS = $commandAst.ToString() + $Env:_CLICK_COMPLETE_WORD_TO_COMPLETE = $wordToComplete + %(prog_name)s | ForEach-Object { + $commandArray = $_ -Split ":::" + $command = $commandArray[0] + $helpString = $commandArray[1] + [System.Management.Automation.CompletionResult]::new( + $command, $command, 'ParameterValue', $helpString) + } + $Env:%(complete_var)s = "" + $Env:_CLICK_COMPLETE_ARGS = "" + $Env:_CLICK_COMPLETE_WORD_TO_COMPLETE = "" +} +Register-ArgumentCompleter -Native -CommandName %(prog_name)s -ScriptBlock $scriptblock +""" + + +class PowerShellComplete(ShellComplete): + name = 'powershell' + source_template = POWERSHELL_COMPLETION_SCRIPT + + def get_completion_args(self) -> tuple[list[str], str]: # pragma: nocover + completion_args = os.getenv('_CLICK_COMPLETE_ARGS', '') + incomplete = os.getenv('_CLICK_COMPLETE_WORD_TO_COMPLETE', '') + cwords = click.parser.split_arg_string(completion_args) + args = cwords[1:] + return args, incomplete + + def format_completion(self, item: click.shell_completion.CompletionItem) -> str: # pragma: nocover + return f'{item.value}:::{item.help or " "}' + + +class PowerCoreComplete(PowerShellComplete): + name = 'pwsh' + + +add_completion_class(PowerShellComplete) +add_completion_class(PowerCoreComplete) + + +def install_powershell(shell: Literal['powershell', 'pwsh']): + # Ok I will explain what I have understood from the algorith I took my inspiration from + # we try to set an execution policy suitable for the current user + subprocess.run([shell, '-Command', 'Set-ExecutionPolicy', 'Unrestricted', '-Scope', 'CurrentUser']) # nosec + + # we get the powershell user profile file where we will store the completion script + try: + result = subprocess.run( # nosec + [shell, '-NoProfile', '-Command', 'echo', '$profile'], check=True, capture_output=True + ) + except subprocess.CalledProcessError: + console.print('[error]Unable to get PowerShell user profile') + raise SystemExit(1) + + user_profile = result.stdout.decode() + user_profile_path = Path(user_profile.strip()) + parent_path = user_profile_path.parent + # we make sure parents directories exist + parent_path.mkdir(parents=True, exist_ok=True) + completion_script = POWERSHELL_COMPLETION_SCRIPT % {'prog_name': 'ws', 'complete_var': '_WS_COMPLETE'} + with open(user_profile_path, 'a') as f: + f.write(f'{completion_script.strip()}\n') def install_bash_zsh(bash: bool = True) -> None: @@ -60,6 +135,8 @@ def _install_completion(shell: str) -> None: install_bash_zsh() elif shell == 'zsh': install_bash_zsh(bash=False) + elif shell in ('powershell', 'pwsh'): + install_powershell(shell) # type: ignore else: install_fish()