diff --git a/lean/commands/__init__.py b/lean/commands/__init__.py index dfa40a41..cbaff0ea 100644 --- a/lean/commands/__init__.py +++ b/lean/commands/__init__.py @@ -11,7 +11,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from os import environ +from copy import copy + +if environ.get("_LEAN_COMPLETE", "").startswith("powershell_"): + from lean.components.util.click_shell_autocomplete import register_shell_autocomplete + register_shell_autocomplete() + from lean.commands.lean import lean +from lean.commands.autocomplete import autocomplete from lean.commands.backtest import backtest from lean.commands.build import build from lean.commands.cloud import cloud @@ -36,6 +44,11 @@ from lean.commands.private_cloud import private_cloud lean.add_command(config) +lean.add_command(autocomplete) +completion = copy(autocomplete) +completion.name = "completion" +completion.hidden = True +lean.add_command(completion) lean.add_command(cloud) lean.add_command(data) lean.add_command(decrypt) diff --git a/lean/commands/autocomplete.py b/lean/commands/autocomplete.py new file mode 100644 index 00000000..31010a2b --- /dev/null +++ b/lean/commands/autocomplete.py @@ -0,0 +1,100 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +from click import Choice, ClickException, Context, echo, group, option, pass_context + +from lean.components.util.click_aliased_command_group import AliasedCommandGroup + + +SHELL_OPTION = option("--shell", + "-s", + type=Choice(["powershell", "bash", "zsh", "fish"], case_sensitive=False), + default=None, + help="Target shell. Auto-detected if not specified.") + + +def _profile_permission_error(exception: PermissionError) -> ClickException: + path = exception.filename or "the shell profile" + return ClickException( + f"Unable to update {path}. " + "The current PowerShell session is still disabled if Lean autocomplete was loaded there. " + "To remove it permanently, close any editor or terminal using that profile, or remove the Lean autocomplete block manually." + ) + + +@group(cls=AliasedCommandGroup, invoke_without_command=True) +@SHELL_OPTION +@pass_context +def autocomplete(ctx: Context, shell: Optional[str]) -> None: + """Print the native shell autocomplete script for your shell. + + \b + PowerShell (current session): + lean autocomplete --shell powershell | Out-String | Invoke-Expression + + \b + Bash or Zsh (current session): + eval "$(lean autocomplete --shell bash)" + + \b + Fish (current session): + lean autocomplete --shell fish | source + """ + if ctx.invoked_subcommand is None: + from lean.components.util.click_shell_autocomplete import get_autocomplete_script + echo(get_autocomplete_script(shell)) + + +@autocomplete.command(name="show", help="Print the native shell autocomplete script for your shell") +@SHELL_OPTION +def show(shell: Optional[str]) -> None: + from lean.components.util.click_shell_autocomplete import get_autocomplete_script + echo(get_autocomplete_script(shell)) + + +@autocomplete.command(name="on", help="Enable shell autocomplete in your shell profile") +@SHELL_OPTION +def on(shell: Optional[str]) -> None: + from lean.components.util.click_shell_autocomplete import install_autocomplete + try: + profile_path = install_autocomplete(shell) + except PermissionError as exception: + raise _profile_permission_error(exception) + + echo(f"Enabled shell autocomplete in {profile_path}") + echo("Open a new terminal session for the change to take effect.") + + +@autocomplete.command(name="off", help="Disable shell autocomplete in your shell profile") +@option("--current-session", is_flag=True, help="Print a script that disables autocomplete in the current shell session.") +@SHELL_OPTION +def off(shell: Optional[str], current_session: bool) -> None: + from lean.components.util.click_shell_autocomplete import get_autocomplete_cleanup_script, uninstall_autocomplete + + if current_session: + echo(get_autocomplete_cleanup_script(shell)) + return + + try: + profile_path, removed = uninstall_autocomplete(shell) + except PermissionError as exception: + raise _profile_permission_error(exception) + + if removed: + echo(f"Disabled shell autocomplete in {profile_path}") + echo("Open a new terminal session for the change to take effect.") + echo("To disable it in this PowerShell session, run `lean autocomplete off` again after reloading the updated script.") + else: + echo(f"Shell autocomplete was not enabled in {profile_path}") diff --git a/lean/commands/cloud/__init__.py b/lean/commands/cloud/__init__.py index 1f13f325..731c7d76 100644 --- a/lean/commands/cloud/__init__.py +++ b/lean/commands/cloud/__init__.py @@ -12,6 +12,7 @@ # limitations under the License. from click import group +from lean.components.util.click_aliased_command_group import AliasedCommandGroup from lean.commands.cloud.backtest import backtest from lean.commands.cloud.live.live import live @@ -21,7 +22,7 @@ from lean.commands.cloud.status import status from lean.commands.cloud.object_store import object_store -@group() +@group(cls=AliasedCommandGroup) def cloud() -> None: """Interact with the QuantConnect cloud.""" # This method is intentionally empty diff --git a/lean/commands/config/__init__.py b/lean/commands/config/__init__.py index c33a6a5c..c4d77d41 100644 --- a/lean/commands/config/__init__.py +++ b/lean/commands/config/__init__.py @@ -12,6 +12,7 @@ # limitations under the License. from click import group +from lean.components.util.click_aliased_command_group import AliasedCommandGroup from lean.commands.config.get import get from lean.commands.config.list import list @@ -19,7 +20,7 @@ from lean.commands.config.unset import unset -@group() +@group(cls=AliasedCommandGroup) def config() -> None: """Configure Lean CLI options.""" # This method is intentionally empty diff --git a/lean/commands/data/__init__.py b/lean/commands/data/__init__.py index a27149db..343a78cf 100644 --- a/lean/commands/data/__init__.py +++ b/lean/commands/data/__init__.py @@ -12,12 +12,13 @@ # limitations under the License. from click import group +from lean.components.util.click_aliased_command_group import AliasedCommandGroup from lean.commands.data.download import download from lean.commands.data.generate import generate -@group() +@group(cls=AliasedCommandGroup) def data() -> None: """Download or generate data for local use.""" # This method is intentionally empty diff --git a/lean/commands/init.py b/lean/commands/init.py index 89dd4a76..450330e8 100644 --- a/lean/commands/init.py +++ b/lean/commands/init.py @@ -207,6 +207,7 @@ def init(organization: Optional[str], language: Optional[str]) -> None: - Synchronizing projects with the cloud: https://www.lean.io/docs/v2/lean-cli/projects/cloud-synchronization Here are some commands to get you going: +- Run `lean autocomplete --shell powershell | Out-String | Invoke-Expression` to enable PowerShell autocomplete in the current session - Run `lean create-project "My Project"` to create a new project with starter code - Run `lean cloud pull` to download all your QuantConnect projects to your local drive - Run `lean backtest "My Project"` to backtest a project locally with the data in {DEFAULT_DATA_DIRECTORY_NAME}/ diff --git a/lean/commands/library/__init__.py b/lean/commands/library/__init__.py index 762ab097..b1711e90 100644 --- a/lean/commands/library/__init__.py +++ b/lean/commands/library/__init__.py @@ -12,12 +12,13 @@ # limitations under the License. from click import group +from lean.components.util.click_aliased_command_group import AliasedCommandGroup from lean.commands.library.add import add from lean.commands.library.remove import remove -@group() +@group(cls=AliasedCommandGroup) def library() -> None: """Manage custom libraries in a project.""" # This method is intentionally empty diff --git a/lean/commands/private_cloud/__init__.py b/lean/commands/private_cloud/__init__.py index b154688c..9ac8b552 100644 --- a/lean/commands/private_cloud/__init__.py +++ b/lean/commands/private_cloud/__init__.py @@ -12,13 +12,14 @@ # limitations under the License. from click import group +from lean.components.util.click_aliased_command_group import AliasedCommandGroup from lean.commands.private_cloud.start import start from lean.commands.private_cloud.stop import stop from lean.commands.private_cloud.add_compute import add_compute -@group() +@group(cls=AliasedCommandGroup) def private_cloud() -> None: """Interact with a QuantConnect private cloud.""" # This method is intentionally empty diff --git a/lean/components/util/click_aliased_command_group.py b/lean/components/util/click_aliased_command_group.py index 68e90cf1..416f4843 100644 --- a/lean/components/util/click_aliased_command_group.py +++ b/lean/components/util/click_aliased_command_group.py @@ -11,35 +11,86 @@ # See the License for the specific language governing permissions and # limitations under the License. -from click import Group +from typing import Any, Callable, Optional, Union, overload + +from click import Command, Context, Group + + +CommandCallback = Callable[..., Any] +CommandDecorator = Callable[[CommandCallback], Command] class AliasedCommandGroup(Group): - """A click.Group wrapper that implements command aliasing.""" + """A click.Group wrapper that implements command aliasing and autocomplete prefix matching.""" + + def get_command(self, ctx: Context, cmd_name: str) -> Optional[Command]: + rv = super().get_command(ctx, cmd_name) + if rv is not None: + return rv + + matches = [] + for name in self.list_commands(ctx): + command = super().get_command(ctx, name) + if command is not None and not command.hidden and name.startswith(cmd_name): + matches.append(name) + + if not matches: + return None + elif len(matches) == 1: + return super().get_command(ctx, matches[0]) + + ctx.fail(f"Too many matches: {', '.join(sorted(matches))}") - def command(self, *args, **kwargs): + @overload + def command(self, __func: CommandCallback) -> Command: + ... + + @overload + def command(self, *args: Any, **kwargs: Any) -> CommandDecorator: + ... + + def command(self, *args: Any, **kwargs: Any) -> Union[CommandDecorator, Command]: aliases = kwargs.pop('aliases', []) - if not args: - cmd_name = kwargs.pop("name", "") - else: - cmd_name = args[0] - args = args[1:] + if not aliases: + return super().command(*args, **kwargs) + + func = None + if args and callable(args[0]): + assert len(args) == 1, "Use 'command(**kwargs)(callable)' to provide arguments." + func = args[0] + args = () - alias_help = f"Alias for '{cmd_name}'" + def _decorator(f: CommandCallback) -> Command: + cmd_kwargs = dict(kwargs) + cmd_name = cmd_kwargs.pop("name", None) + + if args: + if cmd_name is None: + cmd_name = args[0] + cmd_args = args[1:] + else: + cmd_args = args + else: + cmd_name = cmd_name or f.__name__.lower().replace("_", "-") + cmd_args = () + + alias_help = f"Alias for '{cmd_name}'" - def _decorator(f): # Add the main command - cmd = super(AliasedCommandGroup, self).command(name=cmd_name, *args, **kwargs)(f) + cmd = super(AliasedCommandGroup, self).command(*cmd_args, name=cmd_name, **cmd_kwargs)(f) # Add a command to the group for each alias with the same callback but using the alias as name for alias in aliases: alias_cmd = super(AliasedCommandGroup, self).command(name=alias, short_help=alias_help, - *args, - **kwargs)(f) + *cmd_args, + **cmd_kwargs)(f) alias_cmd.params = cmd.params return cmd + if func is not None: + return _decorator(func) + return _decorator diff --git a/lean/components/util/click_group_default_command.py b/lean/components/util/click_group_default_command.py index 7d094d61..38e26d0f 100644 --- a/lean/components/util/click_group_default_command.py +++ b/lean/components/util/click_group_default_command.py @@ -11,9 +11,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from click import Group +from lean.components.util.click_aliased_command_group import AliasedCommandGroup -class DefaultCommandGroup(Group): +class DefaultCommandGroup(AliasedCommandGroup): """allow a default command for a group""" def command(self, *args, **kwargs): diff --git a/lean/components/util/click_shell_autocomplete.py b/lean/components/util/click_shell_autocomplete.py new file mode 100644 index 00000000..437c0eab --- /dev/null +++ b/lean/components/util/click_shell_autocomplete.py @@ -0,0 +1,264 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from pathlib import Path +from platform import system +from typing import Optional + +from click.shell_completion import ShellComplete, add_completion_class, get_completion_class, split_arg_string + +_SOURCE_POWERSHELL = r""" +Register-ArgumentCompleter -Native -CommandName %(prog_name)s -ScriptBlock { + param($wordToComplete, $commandAst, $cursorPosition) + + $Env:%(complete_var)s = "powershell_complete" + $Env:COMP_WORDS = $commandAst.ToString() + $Env:COMP_CWORD = $wordToComplete + + try { + & %(prog_name)s | ForEach-Object { + if ([string]::IsNullOrWhiteSpace($_)) { + return + } + + $item = $_ | ConvertFrom-Json + $tooltip = if ($item.help) { $item.help } else { $item.value } + + [System.Management.Automation.CompletionResult]::new( + $item.value, + $item.value, + 'ParameterValue', + $tooltip + ) + } + } finally { + Remove-Item Env:%(complete_var)s -ErrorAction SilentlyContinue + Remove-Item Env:COMP_WORDS -ErrorAction SilentlyContinue + Remove-Item Env:COMP_CWORD -ErrorAction SilentlyContinue + } +} + +try { + Set-PSReadLineKeyHandler -Key Tab -Function MenuComplete -ErrorAction SilentlyContinue +} catch {} + +function __LeanCliExecutable { + $command = Get-Command %(prog_name)s -CommandType Application -ErrorAction SilentlyContinue | Select-Object -First 1 + if ($command) { + return $command.Source + } + + return "%(prog_name)s" +} + +function lean-autocomplete-off { + Register-ArgumentCompleter -Native -CommandName lean -ScriptBlock { @() } + + try { + Set-PSReadLineOption -PredictionSource None -ErrorAction SilentlyContinue + } catch {} + + Write-Output "Disabled Lean autocomplete for this PowerShell session." + Remove-Item Function:\lean -ErrorAction SilentlyContinue + Remove-Item Function:\lean-autocomplete-on -ErrorAction SilentlyContinue + Remove-Item Function:\lean-autocomplete-off -ErrorAction SilentlyContinue + Remove-Item Function:\__LeanCliExecutable -ErrorAction SilentlyContinue +} + +function lean-autocomplete-on { + $lean = __LeanCliExecutable + & $lean autocomplete --shell powershell | Out-String | Invoke-Expression +} + +function lean { + $lean = __LeanCliExecutable + + if ($args.Count -ge 2 -and ($args[0] -eq "completion" -or $args[0] -eq "autocomplete") -and $args[1] -eq "off") { + lean-autocomplete-off + & $lean @args + return + } + + if ($args.Count -ge 2 -and ($args[0] -eq "completion" -or $args[0] -eq "autocomplete") -and $args[1] -eq "on") { + & $lean @args + lean-autocomplete-on + return + } + + & $lean @args +} +""" + + +class PowerShellComplete(ShellComplete): + """Shell autocomplete for PowerShell.""" + + name = "powershell" + source_template = _SOURCE_POWERSHELL + + def get_completion_args(self) -> tuple[list[str], str]: + cwords = split_arg_string(os.environ.get("COMP_WORDS", "")) + incomplete = os.environ.get("COMP_CWORD", "") + args = cwords[1:] + + if incomplete and args and args[-1] == incomplete: + args.pop() + + return args, incomplete + + def format_completion(self, item) -> str: + return json.dumps({ + "type": item.type, + "value": item.value, + "help": item.help or "" + }, separators=(",", ":")) + + +def register_shell_autocomplete() -> None: + if get_completion_class(PowerShellComplete.name) is None: + add_completion_class(PowerShellComplete) + + +def detect_shell() -> str: + """Auto-detect the current shell environment.""" + if system() == "Windows": + return "powershell" + + shell_path = os.environ.get("SHELL", "/bin/bash") + shell_name = Path(shell_path).name.lower() + + if "zsh" in shell_name: + return "zsh" + + if "fish" in shell_name: + return "fish" + + return "bash" + + +def get_autocomplete_script(shell: Optional[str], prog_name: str = "lean") -> str: + register_shell_autocomplete() + + shell_name = (shell or detect_shell()).lower() + complete_var = f"_{prog_name.replace('-', '_').replace('.', '_')}_COMPLETE".upper() + completion_class = get_completion_class(shell_name) + + if completion_class is None: + supported_shells = ", ".join(sorted(["bash", "fish", "powershell", "zsh"])) + raise RuntimeError(f"Unsupported shell '{shell_name}'. Supported shells: {supported_shells}") + + return completion_class(None, {}, prog_name, complete_var).source() + + +def get_autocomplete_cleanup_script(shell: Optional[str], prog_name: str = "lean") -> str: + shell_name = (shell or detect_shell()).lower() + + if shell_name == "powershell": + return f""" +Register-ArgumentCompleter -Native -CommandName {prog_name} -ScriptBlock {{ @() }} + +try {{ + Set-PSReadLineOption -PredictionSource None -ErrorAction SilentlyContinue +}} catch {{}} + +Remove-Item Function:\\lean -ErrorAction SilentlyContinue +Remove-Item Function:\\lean-autocomplete-on -ErrorAction SilentlyContinue +Remove-Item Function:\\lean-autocomplete-off -ErrorAction SilentlyContinue +Remove-Item Function:\\__LeanCliExecutable -ErrorAction SilentlyContinue + +function lean-autocomplete-on {{ + & {prog_name} autocomplete --shell powershell | Out-String | Invoke-Expression +}} +""".strip() + + return "" + + +def get_profile_path(shell: Optional[str]) -> Path: + shell_name = (shell or detect_shell()).lower() + + if shell_name == "powershell": + if system() == "Windows": + return Path.home() / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + + return Path.home() / ".config" / "powershell" / "Microsoft.PowerShell_profile.ps1" + + if shell_name == "zsh": + return Path.home() / ".zshrc" + + if shell_name == "fish": + return Path.home() / ".config" / "fish" / "completions" / "lean.fish" + + return Path.home() / ".bashrc" + + +def install_autocomplete(shell: Optional[str], prog_name: str = "lean") -> Path: + profile_path = get_profile_path(shell) + marker_start = "# >>> lean autocomplete >>>" + marker_end = "# <<< lean autocomplete <<<" + script = get_autocomplete_script(shell, prog_name).strip() + + content = profile_path.read_text(encoding="utf-8") if profile_path.exists() else "" + if marker_start in content: + return profile_path + + profile_path.parent.mkdir(parents=True, exist_ok=True) + block = f"\n{marker_start}\n{script}\n{marker_end}\n" + + with profile_path.open("a", encoding="utf-8") as file: + file.write(block) + + return profile_path + + +def uninstall_autocomplete(shell: Optional[str]) -> tuple[Path, bool]: + profile_path = get_profile_path(shell) + + if not profile_path.exists(): + return profile_path, False + + content = profile_path.read_text(encoding="utf-8") + markers = [ + ("# >>> lean autocomplete >>>", "# <<< lean autocomplete <<<"), + ("# >>> lean completion >>>", "# <<< lean completion <<<") + ] + + start_index = -1 + end_index = -1 + for marker_start, marker_end in markers: + start_index = content.find(marker_start) + if start_index == -1: + continue + + end_index = content.find(marker_end, start_index) + if end_index != -1: + break + + if start_index == -1 or end_index == -1: + return profile_path, False + + end_index += len(marker_end) + new_content = content[:start_index].rstrip("\n") + tail = content[end_index:].lstrip("\n") + + if new_content and tail: + new_content = f"{new_content}\n{tail}" + elif tail: + new_content = tail + elif new_content: + new_content = f"{new_content}\n" + + profile_path.write_text(new_content, encoding="utf-8") + return profile_path, True diff --git a/lean/main.py b/lean/main.py index 061c721b..2ee9044d 100644 --- a/lean/main.py +++ b/lean/main.py @@ -90,12 +90,16 @@ def _ensure_win32_available() -> None: def main() -> None: """This function is the entrypoint when running a Lean command in a terminal.""" + from click.exceptions import Exit + try: lean.main(standalone_mode=False) temp_manager = container.temp_manager if temp_manager.delete_temporary_directories_when_done: temp_manager.delete_temporary_directories() + except Exit as e: + exit(e.exit_code) except Exception as exception: from traceback import format_exc from click import UsageError, Abort diff --git a/tests/test_autocomplete.py b/tests/test_autocomplete.py new file mode 100644 index 00000000..f7aadb00 --- /dev/null +++ b/tests/test_autocomplete.py @@ -0,0 +1,205 @@ +# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals. +# Lean CLI v1.0. Copyright 2021 QuantConnect Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from tempfile import TemporaryDirectory +from pathlib import Path +from unittest.mock import patch + +from click.testing import CliRunner + +from lean.commands import lean +from lean.components.util.click_shell_autocomplete import register_shell_autocomplete + + +def test_hidden_completion_alias_prints_powershell_script() -> None: + result = CliRunner().invoke(lean, ["completion", "--shell", "powershell"]) + + assert result.exit_code == 0 + assert "Register-ArgumentCompleter -Native -CommandName lean" in result.output + assert "_LEAN_COMPLETE" in result.output + assert "Set-PSReadLineOption -PredictionSource HistoryAndPlugin" not in result.output + assert "Set-PSReadLineKeyHandler -Key Tab -Function MenuComplete" in result.output + assert "function lean-autocomplete-off" in result.output + assert "function lean-autocomplete-on" in result.output + assert "function lean {" in result.output + assert '__LeanCliExecutable' in result.output + assert result.output.index("lean-autocomplete-off") < result.output.index("& $lean @args") + + +def test_autocomplete_command_prints_powershell_script() -> None: + result = CliRunner().invoke(lean, ["autocomplete", "--shell", "powershell"]) + + assert result.exit_code == 0 + assert "Register-ArgumentCompleter -Native -CommandName lean" in result.output + assert "& $lean autocomplete --shell powershell | Out-String | Invoke-Expression" in result.output + + +def test_lean_help_shows_autocomplete_not_hidden_completion_alias() -> None: + result = CliRunner().invoke(lean, ["--help"]) + + assert result.exit_code == 0 + assert "autocomplete" in result.output + assert "completion" not in result.output + + +def test_click_shell_autocomplete_prints_powershell_source_script() -> None: + register_shell_autocomplete() + + result = CliRunner().invoke(lean, [], prog_name="lean", env={ + "_LEAN_COMPLETE": "powershell_source" + }) + + assert result.exit_code == 0 + assert "Register-ArgumentCompleter -Native -CommandName lean" in result.output + + +def test_click_shell_autocomplete_returns_powershell_completions() -> None: + register_shell_autocomplete() + + result = CliRunner().invoke(lean, [], prog_name="lean", env={ + "_LEAN_COMPLETE": "powershell_complete", + "COMP_WORDS": "lean cl", + "COMP_CWORD": "cl" + }) + + assert result.exit_code == 0 + + completions = [json.loads(line) for line in result.output.strip().splitlines()] + completion_values = [item["value"] for item in completions] + assert "cloud" in completion_values + + +def test_click_shell_autocomplete_prints_bash_source_script() -> None: + result = CliRunner().invoke(lean, [], prog_name="lean", env={ + "_LEAN_COMPLETE": "bash_source" + }) + + assert result.exit_code == 0 + assert "complete -o nosort -F" in result.output + + +def test_hidden_completion_alias_on_writes_powershell_profile() -> None: + with TemporaryDirectory() as directory, patch.object(Path, "home", return_value=Path(directory)): + result = CliRunner().invoke(lean, ["completion", "on", "--shell", "powershell"]) + + assert result.exit_code == 0 + + profile_path = Path(directory) / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + assert profile_path.exists() + + content = profile_path.read_text(encoding="utf-8") + assert "# >>> lean autocomplete >>>" in content + assert "Register-ArgumentCompleter -Native -CommandName lean" in content + + +def test_autocomplete_on_writes_powershell_profile() -> None: + with TemporaryDirectory() as directory, patch.object(Path, "home", return_value=Path(directory)): + result = CliRunner().invoke(lean, ["autocomplete", "on", "--shell", "powershell"]) + + assert result.exit_code == 0 + + profile_path = Path(directory) / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + assert profile_path.exists() + + content = profile_path.read_text(encoding="utf-8") + assert "# >>> lean autocomplete >>>" in content + assert "Register-ArgumentCompleter -Native -CommandName lean" in content + + +def test_autocomplete_on_auto_detects_powershell_profile() -> None: + with TemporaryDirectory() as directory, \ + patch.object(Path, "home", return_value=Path(directory)), \ + patch("lean.components.util.click_shell_autocomplete.system", return_value="Windows"): + result = CliRunner().invoke(lean, ["autocomplete", "on"]) + + assert result.exit_code == 0 + + profile_path = Path(directory) / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + assert profile_path.exists() + assert "# >>> lean autocomplete >>>" in profile_path.read_text(encoding="utf-8") + + +def test_hidden_completion_alias_off_removes_legacy_powershell_profile_block() -> None: + with TemporaryDirectory() as directory, patch.object(Path, "home", return_value=Path(directory)): + profile_path = Path(directory) / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + profile_path.parent.mkdir(parents=True, exist_ok=True) + profile_path.write_text( + "# before\n# >>> lean completion >>>\nlean block\n# <<< lean completion <<<\n# after\n", + encoding="utf-8" + ) + + result = CliRunner().invoke(lean, ["completion", "off", "--shell", "powershell"]) + + assert result.exit_code == 0 + + content = profile_path.read_text(encoding="utf-8") + assert "# >>> lean completion >>>" not in content + assert "# before" in content + assert "# after" in content + assert "lean autocomplete off" in result.output + + +def test_autocomplete_off_removes_powershell_profile_block() -> None: + with TemporaryDirectory() as directory, patch.object(Path, "home", return_value=Path(directory)): + profile_path = Path(directory) / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + profile_path.parent.mkdir(parents=True, exist_ok=True) + profile_path.write_text( + "# before\n# >>> lean autocomplete >>>\nlean block\n# <<< lean autocomplete <<<\n# after\n", + encoding="utf-8" + ) + + result = CliRunner().invoke(lean, ["autocomplete", "off", "--shell", "powershell"]) + + assert result.exit_code == 0 + assert "# >>> lean autocomplete >>>" not in profile_path.read_text(encoding="utf-8") + assert "lean autocomplete off" in result.output + + +def test_autocomplete_off_auto_detects_powershell_profile() -> None: + with TemporaryDirectory() as directory, \ + patch.object(Path, "home", return_value=Path(directory)), \ + patch("lean.components.util.click_shell_autocomplete.system", return_value="Windows"): + profile_path = Path(directory) / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + profile_path.parent.mkdir(parents=True, exist_ok=True) + profile_path.write_text( + "# before\n# >>> lean autocomplete >>>\nlean block\n# <<< lean autocomplete <<<\n# after\n", + encoding="utf-8" + ) + + result = CliRunner().invoke(lean, ["autocomplete", "off"]) + + assert result.exit_code == 0 + assert "# >>> lean autocomplete >>>" not in profile_path.read_text(encoding="utf-8") + assert "lean autocomplete off" in result.output + + +def test_autocomplete_off_current_session_prints_powershell_cleanup_script() -> None: + result = CliRunner().invoke(lean, ["autocomplete", "off", "--shell", "powershell", "--current-session"]) + + assert result.exit_code == 0 + assert "Register-ArgumentCompleter -Native -CommandName lean -ScriptBlock { @() }" in result.output + assert "Set-PSReadLineOption -PredictionSource None" in result.output + assert "Remove-Item Function:\\lean" in result.output + assert "function lean-autocomplete-on" in result.output + assert "autocomplete --shell powershell" in result.output + + +def test_autocomplete_off_shows_clear_error_when_profile_cannot_be_updated() -> None: + with patch("lean.components.util.click_shell_autocomplete.uninstall_autocomplete", + side_effect=PermissionError(13, "Permission denied", "profile.ps1")): + result = CliRunner().invoke(lean, ["autocomplete", "off", "--shell", "powershell"]) + + assert result.exit_code != 0 + assert "Unable to update profile.ps1" in result.output + assert "The current PowerShell session is still disabled" in result.output diff --git a/tests/test_click_aliased_command_group.py b/tests/test_click_aliased_command_group.py index 33c62cd8..569bc77b 100644 --- a/tests/test_click_aliased_command_group.py +++ b/tests/test_click_aliased_command_group.py @@ -80,3 +80,59 @@ def command() -> None: assert len(aliases_help) == len(aliases_help) assert all(f"Alias for '{command_name}'" in alias_help for alias_help in aliases_help) assert main_command_doc in main_command_help + + +def test_aliased_command_group_resolves_unique_prefix_match() -> None: + @click.group(cls=AliasedCommandGroup) + def group() -> None: + pass + + @group.command() + def cloud() -> None: + click.echo("cloud") + + result = CliRunner().invoke(group, ["cl"]) + + assert result.exit_code == 0 + assert result.output == "cloud\n" + + +def test_aliased_command_group_fails_when_prefix_is_ambiguous() -> None: + @click.group(cls=AliasedCommandGroup) + def group() -> None: + pass + + @group.command() + def cloud() -> None: + pass + + @group.command() + def config() -> None: + pass + + result = CliRunner().invoke(group, ["c"]) + + assert result.exit_code != 0 + assert "Too many matches: cloud, config" in result.output + + +def test_aliased_command_group_ignores_hidden_commands_for_prefix_matching() -> None: + @click.group(cls=AliasedCommandGroup) + def group() -> None: + pass + + @group.command(hidden=True) + def completion() -> None: + click.echo("completion") + + @group.command() + def cloud() -> None: + click.echo("cloud") + + prefix_result = CliRunner().invoke(group, ["c"]) + exact_result = CliRunner().invoke(group, ["completion"]) + + assert prefix_result.exit_code == 0 + assert prefix_result.output == "cloud\n" + assert exact_result.exit_code == 0 + assert exact_result.output == "completion\n" diff --git a/tests/test_main.py b/tests/test_main.py index ffea5974..9acef3f5 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -35,3 +35,26 @@ def test_lean_shows_error_when_running_unknown_command() -> None: assert result.exit_code != 0 assert "No such command" in result.output + + +def test_lean_runs_top_level_commands_by_unique_prefix() -> None: + result = CliRunner().invoke(lean, ["cl", "--help"]) + + assert result.exit_code == 0 + assert "Interact with the QuantConnect cloud." in result.output + assert "backtest" in result.output + + +def test_lean_runs_nested_commands_by_unique_prefix() -> None: + result = CliRunner().invoke(lean, ["cloud", "st", "--help"]) + + assert result.exit_code == 0 + assert "Show the live trading status of a project in the cloud." in result.output + assert "PROJECT" in result.output + + +def test_lean_reports_ambiguous_prefixes() -> None: + result = CliRunner().invoke(lean, ["c"]) + + assert result.exit_code != 0 + assert "Too many matches: cloud, config, create-project" in result.output