feat(completion): implemented powershell cli completion

tests/test_completion.py → 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')

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]:
+        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:
+        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()