diff --git a/opnsense_cli/cli.py b/opnsense_cli/cli.py index 904cc35..bb9539e 100755 --- a/opnsense_cli/cli.py +++ b/opnsense_cli/cli.py @@ -5,7 +5,7 @@ from opnsense_cli.click_addons.callbacks import defaults_from_configfile, expand_path, get_default_config_dir from opnsense_cli.api.client import ApiClient from opnsense_cli.click_addons.command_autoloader import ClickCommandAutoloader - +from opnsense_cli.click_addons.command_tree import _build_command_tree, _print_tree CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"]) @@ -132,13 +132,13 @@ def cli(ctx, **kwargs): kwargs["timeout"], ) - autoloader = ClickCommandAutoloader(cli) autoloader.autoload("opnsense_cli.commands.core") autoloader.autoload("opnsense_cli.commands.plugin") autoloader.autoload("opnsense_cli.commands.new") autoloader.autoload("opnsense_cli.commands.completion") autoloader.autoload("opnsense_cli.commands.version") +autoloader.autoload("opnsense_cli.commands.tree") if __name__ == "__main__": cli() diff --git a/opnsense_cli/click_addons/command_autoloader.py b/opnsense_cli/click_addons/command_autoloader.py index b3f0821..5b767f0 100644 --- a/opnsense_cli/click_addons/command_autoloader.py +++ b/opnsense_cli/click_addons/command_autoloader.py @@ -2,11 +2,11 @@ import importlib import importlib.util - class ClickCommandAutoloader: def __init__(self, click_core_group): - self.loaded_modules = set() + self.loaded_modules = [] self.loaded_classes = [] + self.subgroups = {} self.click_core_group = click_core_group def autoload(self, module_name=None, ignore_dirs=None): @@ -37,36 +37,34 @@ def autoload(self, module_name=None, ignore_dirs=None): path = spec.submodule_search_locations[0] (root_dir, command_group_dirs, files) = list(os.walk(path))[0] - command_group_dirs = [directory for directory in command_group_dirs if directory not in ignore_dirs] if not command_group_dirs: path, file = os.path.split(root_dir) command_group_dirs = [file] module_path_components = module_name.split(".") - module_name = ".".join(module_path_components[0 : len(module_path_components) - 1]) + module_name = ".".join(module_path_components[:-1]) for command_group_dir in command_group_dirs: - command_group_files = list(os.walk(f"{path}/{command_group_dir}"))[0][2] + command_group_files = sorted(list(os.walk(f"{path}/{command_group_dir}"))[0][2]) for command_group_file in command_group_files: - import_name = f"{module_name}.{command_group_dir}" - class_name = f"{command_group_dir}" - - if command_group_file != "__init__.py": + if command_group_file == "__init__.py": + import_name = f"{module_name}.{command_group_dir}" + class_name = command_group_dir + else: _subname_ = os.path.splitext(command_group_file)[0] import_name = f"{module_name}.{command_group_dir}.{_subname_}" - class_name = f"{_subname_}" + class_name = _subname_ module = importlib.import_module(import_name) click_group = getattr(module, class_name) - self.loaded_modules.add(module) + self.loaded_modules.append(module) self.loaded_classes.append(click_group) if command_group_file == "__init__.py": self.click_core_group.add_command(click_group) - else: - click_group.add_command(click_group) + self.subgroups[command_group_dir] = click_group - return click_group + return self.click_core_group diff --git a/opnsense_cli/click_addons/command_tree.py b/opnsense_cli/click_addons/command_tree.py new file mode 100644 index 0000000..b1319d2 --- /dev/null +++ b/opnsense_cli/click_addons/command_tree.py @@ -0,0 +1,43 @@ +import click + + +class _CommandWrapper(object): + def __init__(self, command=None, children=None): + self.command = command + self.children = [] + + @property + def name(self): + return self.command.name + + +def _build_command_tree(click_command): + wrapper = _CommandWrapper(click_command) + + if isinstance(click_command, click.core.Group): + for _, cmd in click_command.commands.items(): + if not getattr(cmd, "hidden", False): + wrapper.children.append(_build_command_tree(cmd)) + + return wrapper + + +def _print_tree(command, depth=0, is_last_item=False, is_last_parent=False): + if depth == 0: + prefix = '' + tree_item = '' + else: + prefix = ' ' if is_last_parent else '│ ' + tree_item = '└── ' if is_last_item else '├── ' + + root_line = '│ ' if is_last_parent else '' + + line = root_line + prefix * (depth - 1) + tree_item + command.name + + click.echo(line) + + for i, child in enumerate(sorted(command.children, key=lambda x: x.name)): + _print_tree(child, + depth=(depth + 1), + is_last_item=(i == (len(command.children) - 1)), + is_last_parent=is_last_item) diff --git a/opnsense_cli/commands/tree/__init__.py b/opnsense_cli/commands/tree/__init__.py new file mode 100644 index 0000000..76b8429 --- /dev/null +++ b/opnsense_cli/commands/tree/__init__.py @@ -0,0 +1,14 @@ +import click +from opnsense_cli.click_addons.command_tree import _build_command_tree, _print_tree + + +@click.command() +@click.pass_context +def tree(ctx): + """ + Show the command tree of your CLI + """ + root_cmd = _build_command_tree(ctx.find_root().command) + _print_tree(root_cmd) + + diff --git a/opnsense_cli/commands/tree/tests/__init__.py b/opnsense_cli/commands/tree/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/opnsense_cli/commands/tree/tests/test_tree.py b/opnsense_cli/commands/tree/tests/test_tree.py new file mode 100644 index 0000000..5302bd6 --- /dev/null +++ b/opnsense_cli/commands/tree/tests/test_tree.py @@ -0,0 +1,56 @@ +import unittest +from click.testing import CliRunner +from opnsense_cli.commands.tree import tree +import click +import textwrap + + +class TestTreeCommands(unittest.TestCase): + def test_tree(self): + @click.group(name='root') + def root(): + pass + + @root.command(name='command-one') + def command_one(): + pass + + @root.command(name='command-two') + def command_two(): + pass + + @click.group(name='sub_level1') + def sub_level1(): + pass + + @click.group(name='sub_level2') + def sub_level2(): + pass + + root.add_command(tree) + + root.add_command(sub_level1) + sub_level1.add_command(command_one) + sub_level1.add_command(command_two) + + sub_level1.add_command(sub_level2) + sub_level2.add_command(command_one) + sub_level2.add_command(command_two) + + runner = CliRunner() + result = runner.invoke(root, ['tree']) + + tree_output = textwrap.dedent("""\ + root + ├── command-one + ├── command-two + ├── sub_level1 + │ ├── command-one + │ ├── command-two + │ └── sub_level2 + │ ├── command-one + │ └── command-two + └── tree + """) + + self.assertEqual(tree_output, result.output)