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/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..5950dbf3 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 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/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