diff --git a/lean/commands/__init__.py b/lean/commands/__init__.py index dfa40a41..c5d5b4d3 100644 --- a/lean/commands/__init__.py +++ b/lean/commands/__init__.py @@ -12,6 +12,7 @@ # limitations under the License. from lean.commands.lean import lean +from lean.commands.completion import completion from lean.commands.backtest import backtest from lean.commands.build import build from lean.commands.cloud import cloud @@ -36,6 +37,7 @@ from lean.commands.private_cloud import private_cloud lean.add_command(config) +lean.add_command(completion) lean.add_command(cloud) lean.add_command(data) lean.add_command(decrypt) 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/completion.py b/lean/commands/completion.py new file mode 100644 index 00000000..418001bf --- /dev/null +++ b/lean/commands/completion.py @@ -0,0 +1,74 @@ +# 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, Context, echo, group, option, pass_context + +from lean.components.util.click_aliased_command_group import AliasedCommandGroup +from lean.components.util.click_shell_completion import get_completion_script, install_completion, uninstall_completion + + +SHELL_OPTION = option("--shell", + "-s", + type=Choice(["powershell", "bash", "zsh", "fish"], case_sensitive=False), + default=None, + help="Target shell. Auto-detected if not specified.") + + +@group(cls=AliasedCommandGroup, invoke_without_command=True) +@SHELL_OPTION +@pass_context +def completion(ctx: Context, shell: Optional[str]) -> None: + """Print the native shell completion script for your shell. + + \b + PowerShell (current session): + lean completion --shell powershell | Out-String | Invoke-Expression + + \b + Bash or Zsh (current session): + eval "$(lean completion --shell bash)" + + \b + Fish (current session): + lean completion --shell fish | source + """ + if ctx.invoked_subcommand is None: + echo(get_completion_script(shell)) + + +@completion.command(name="show", help="Print the native shell completion script for your shell") +@SHELL_OPTION +def show(shell: Optional[str]) -> None: + echo(get_completion_script(shell)) + + +@completion.command(name="on", help="Enable shell completion in your shell profile") +@SHELL_OPTION +def on(shell: Optional[str]) -> None: + profile_path = install_completion(shell) + echo(f"Enabled shell completion in {profile_path}") + echo("Open a new terminal session for the change to take effect.") + + +@completion.command(name="off", help="Disable shell completion in your shell profile") +@SHELL_OPTION +def off(shell: Optional[str]) -> None: + profile_path, removed = uninstall_completion(shell) + + if removed: + echo(f"Disabled shell completion in {profile_path}") + echo("Open a new terminal session for the change to take effect.") + else: + echo(f"Shell completion was not enabled in {profile_path}") 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..87ea1ea0 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 completion --shell powershell | Out-String | Invoke-Expression` to enable PowerShell completion 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/lean.py b/lean/commands/lean.py index 749000bb..eb1ecc6d 100644 --- a/lean/commands/lean.py +++ b/lean/commands/lean.py @@ -16,9 +16,12 @@ from lean import __version__ from lean.click import verbose_option from lean.components.util.click_aliased_command_group import AliasedCommandGroup +from lean.components.util.click_shell_completion import register_shell_completion from lean.container import container from lean.models.errors import MoreInfoError +register_shell_completion() + @group(cls=AliasedCommandGroup, invoke_without_command=True) @option("--version", is_flag=True, is_eager=True, help="Show the version and exit.") 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..3267fbef 100644 --- a/lean/components/util/click_aliased_command_group.py +++ b/lean/components/util/click_aliased_command_group.py @@ -15,29 +15,49 @@ class AliasedCommandGroup(Group): - """A click.Group wrapper that implements command aliasing.""" + """A click.Group wrapper that implements command aliasing and auto-completion/prefix matching.""" + + def get_command(self, ctx, cmd_name): + rv = super().get_command(ctx, cmd_name) + if rv is not None: + return rv + + matches = [x for x in self.list_commands(ctx) if x.startswith(cmd_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): aliases = kwargs.pop('aliases', []) - if not args: - cmd_name = kwargs.pop("name", "") - else: - cmd_name = args[0] - args = args[1:] - - alias_help = f"Alias for '{cmd_name}'" + if not aliases: + return super().command(*args, **kwargs) def _decorator(f): + if args: + cmd_name = args[0] + cmd_args = args[1:] + else: + cmd_name = kwargs.get("name", f.__name__.lower().replace("_", "-")) + cmd_args = () + + alias_help = f"Alias for '{cmd_name}'" + cmd_kwargs = dict(kwargs) + cmd_kwargs.pop("name", None) + # 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 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_completion.py b/lean/components/util/click_shell_completion.py new file mode 100644 index 00000000..be963a1a --- /dev/null +++ b/lean/components/util/click_shell_completion.py @@ -0,0 +1,181 @@ +# 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 + } +} +""" + + +class PowerShellComplete(ShellComplete): + """Shell completion 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_completion() -> 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_completion_script(shell: Optional[str], prog_name: str = "lean") -> str: + register_shell_completion() + + 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_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_completion(shell: Optional[str], prog_name: str = "lean") -> Path: + profile_path = get_profile_path(shell) + marker_start = "# >>> lean completion >>>" + marker_end = "# <<< lean completion <<<" + script = get_completion_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_completion(shell: Optional[str]) -> tuple[Path, bool]: + profile_path = get_profile_path(shell) + marker_start = "# >>> lean completion >>>" + marker_end = "# <<< lean completion <<<" + + if not profile_path.exists(): + return profile_path, False + + content = profile_path.read_text(encoding="utf-8") + start_index = content.find(marker_start) + if start_index == -1: + return profile_path, False + + end_index = content.find(marker_end, start_index) + if 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_click_aliased_command_group.py b/tests/test_click_aliased_command_group.py index 33c62cd8..023a7da8 100644 --- a/tests/test_click_aliased_command_group.py +++ b/tests/test_click_aliased_command_group.py @@ -80,3 +80,37 @@ 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 diff --git a/tests/test_completion.py b/tests/test_completion.py new file mode 100644 index 00000000..9f32a53e --- /dev/null +++ b/tests/test_completion.py @@ -0,0 +1,90 @@ +# 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 pathlib import Path + +from click.testing import CliRunner + +from lean.commands import lean + + +def test_completion_command_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 + + +def test_click_shell_completion_prints_powershell_source_script() -> None: + 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_completion_returns_powershell_completions() -> None: + 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_completion_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_completion_on_writes_powershell_profile() -> None: + result = CliRunner().invoke(lean, ["completion", "on", "--shell", "powershell"]) + + assert result.exit_code == 0 + + profile_path = Path.home() / "Documents" / "PowerShell" / "Microsoft.PowerShell_profile.ps1" + assert profile_path.exists() + + content = profile_path.read_text(encoding="utf-8") + assert "# >>> lean completion >>>" in content + assert "Register-ArgumentCompleter -Native -CommandName lean" in content + + +def test_completion_off_removes_powershell_profile_block() -> None: + profile_path = Path.home() / "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 diff --git a/tests/test_main.py b/tests/test_main.py index ffea5974..da206bef 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, completion, config, create-project" in result.output