diff --git a/src/shelloracle/bootstrap.py b/src/shelloracle/bootstrap.py index 95fd746..1266187 100644 --- a/src/shelloracle/bootstrap.py +++ b/src/shelloracle/bootstrap.py @@ -1,6 +1,7 @@ from __future__ import annotations import inspect +import platform import shutil from pathlib import Path from typing import TYPE_CHECKING, Any @@ -37,7 +38,7 @@ def replace_home_with_tilde(path: Path) -> Path: return Path("~") / relative_path -supported_shells = ("zsh", "bash", "fish") +supported_shells = ("zsh", "bash", "fish", "pwsh") def get_installed_shells() -> list[str]: @@ -50,6 +51,8 @@ def get_bundled_script_path(shell: str) -> Path: return shell_dir / "shelloracle.zsh" if shell == "fish": return shell_dir / "shelloracle.fish" + if shell == "pwsh": + return shell_dir / "shelloracle.ps1" return shell_dir / "shelloracle.bash" @@ -58,6 +61,8 @@ def get_script_path(shell: str) -> Path: return Path.home() / ".shelloracle.zsh" if shell == "fish": return Path.home() / ".shelloracle.fish" + if shell == "pwsh": + return Path.home() / ".shelloracle.ps1" return Path.home() / ".shelloracle.bash" @@ -66,6 +71,10 @@ def get_rc_path(shell: str) -> Path: return Path.home() / ".zshrc" if shell == "fish": return Path.home() / ".config/fish/config.fish" + if shell == "pwsh": + if platform.system() == "Windows": + return Path.home() / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + return Path.home() / ".config" / "powershell" / "Microsoft.PowerShell_profile.ps1" return Path.home() / ".bashrc" @@ -78,11 +87,15 @@ def write_script_home(shell: str) -> None: def update_rc(shell: str) -> None: rc_path = get_rc_path(shell) + rc_path.parent.mkdir(parents=True, exist_ok=True) rc_path.touch(exist_ok=True) with rc_path.open("r") as file: rc_content = file.read() if shell == "fish": line = f"if test -f {get_script_path(shell)}; source {get_script_path(shell)}; end" + elif shell == "pwsh": + shelloracle_script = get_script_path(shell) + line = f". {shelloracle_script}" else: shelloracle_script = get_script_path(shell) line = f"[ -f {shelloracle_script} ] && source {shelloracle_script}" diff --git a/src/shelloracle/providers/__init__.py b/src/shelloracle/providers/__init__.py index 57e446d..101c383 100644 --- a/src/shelloracle/providers/__init__.py +++ b/src/shelloracle/providers/__init__.py @@ -1,6 +1,7 @@ from __future__ import annotations import abc +import os from abc import abstractmethod from typing import TYPE_CHECKING, Generic, TypeVar @@ -9,9 +10,9 @@ from shelloracle.config import Configuration -system_prompt = ( +_SYSTEM_PROMPT_TEMPLATE = ( "Based on the following user description, generate a corresponding shell command. Focus solely " - "on interpreting the requirements and translating them into a single, executable Bash command. " + "on interpreting the requirements and translating them into a single, executable {shell} command. " "Ensure accuracy and relevance to the user's description. The output should be a valid shell " "command that directly aligns with the user's intent, ready for execution in a command-line " "environment. Do not output anything except for the command. No code block, no English explanation, " @@ -19,6 +20,22 @@ ) +def get_system_prompt() -> str: + """Return the system prompt, adjusted for the active shell. + + The shell integration scripts set the SHOR_SHELL environment variable so + that the generated command matches the syntax of the calling shell. + + :return: system prompt string + """ + shell = os.environ.get("SHOR_SHELL", "Bash") + shell_name = "PowerShell" if shell.lower() == "powershell" else "Bash" + return _SYSTEM_PROMPT_TEMPLATE.format(shell=shell_name) + + +system_prompt = get_system_prompt() + + class ProviderError(Exception): """LLM providers raise this error to gracefully indicate something has gone wrong.""" diff --git a/src/shelloracle/shell/shelloracle.ps1 b/src/shelloracle/shell/shelloracle.ps1 new file mode 100644 index 0000000..32e1c7d --- /dev/null +++ b/src/shelloracle/shell/shelloracle.ps1 @@ -0,0 +1,14 @@ +Set-PSReadLineKeyHandler -Key Ctrl+f -ScriptBlock { + $line = $null + $cursor = $null + [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$line, [ref]$cursor) + $env:SHOR_DEFAULT_PROMPT = $line + $env:SHOR_SHELL = "powershell" + $output = & shor + $env:SHOR_DEFAULT_PROMPT = $null + $env:SHOR_SHELL = $null + if ($LASTEXITCODE -eq 0) { + [Microsoft.PowerShell.PSConsoleReadLine]::ReplaceLine($output) + [Microsoft.PowerShell.PSConsoleReadLine]::EndOfLine() + } +} diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py new file mode 100644 index 0000000..bcb209b --- /dev/null +++ b/tests/test_bootstrap.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import platform +from unittest.mock import patch + +from shelloracle.bootstrap import ( + get_bundled_script_path, + get_rc_path, + get_script_path, + supported_shells, +) +from shelloracle.providers import get_system_prompt + + +class TestSupportedShells: + def test_pwsh_in_supported_shells(self): + assert "pwsh" in supported_shells + + +class TestGetBundledScriptPath: + def test_zsh(self): + path = get_bundled_script_path("zsh") + assert path.name == "shelloracle.zsh" + + def test_fish(self): + path = get_bundled_script_path("fish") + assert path.name == "shelloracle.fish" + + def test_bash(self): + path = get_bundled_script_path("bash") + assert path.name == "shelloracle.bash" + + def test_pwsh(self): + path = get_bundled_script_path("pwsh") + assert path.name == "shelloracle.ps1" + assert path.exists() + + +class TestGetScriptPath: + def test_zsh(self): + path = get_script_path("zsh") + assert path.name == ".shelloracle.zsh" + + def test_fish(self): + path = get_script_path("fish") + assert path.name == ".shelloracle.fish" + + def test_bash(self): + path = get_script_path("bash") + assert path.name == ".shelloracle.bash" + + def test_pwsh(self): + path = get_script_path("pwsh") + assert path.name == ".shelloracle.ps1" + + +class TestGetRcPath: + def test_zsh(self): + path = get_rc_path("zsh") + assert path.name == ".zshrc" + + def test_fish(self): + path = get_rc_path("fish") + assert path.name == "config.fish" + + def test_bash(self): + path = get_rc_path("bash") + assert path.name == ".bashrc" + + def test_pwsh_windows(self): + with patch.object(platform, "system", return_value="Windows"): + path = get_rc_path("pwsh") + assert path.name == "Microsoft.PowerShell_profile.ps1" + assert "Documents" in path.parts + + def test_pwsh_non_windows(self): + with patch.object(platform, "system", return_value="Darwin"): + path = get_rc_path("pwsh") + assert path.name == "Microsoft.PowerShell_profile.ps1" + assert ".config" in path.parts + + +class TestGetSystemPrompt: + def test_default_is_bash(self, monkeypatch): + monkeypatch.delenv("SHOR_SHELL", raising=False) + prompt = get_system_prompt() + assert "Bash" in prompt + assert "PowerShell" not in prompt + + def test_powershell(self, monkeypatch): + monkeypatch.setenv("SHOR_SHELL", "powershell") + prompt = get_system_prompt() + assert "PowerShell" in prompt + assert "Bash" not in prompt + + def test_powershell_case_insensitive(self, monkeypatch): + monkeypatch.setenv("SHOR_SHELL", "POWERSHELL") + prompt = get_system_prompt() + assert "PowerShell" in prompt + + def test_unknown_shell_falls_back_to_bash(self, monkeypatch): + monkeypatch.setenv("SHOR_SHELL", "unknown-shell") + prompt = get_system_prompt() + assert "Bash" in prompt