diff --git a/software/script/chameleon_cli_main.py b/software/script/chameleon_cli_main.py index d898479..5212ac2 100755 --- a/software/script/chameleon_cli_main.py +++ b/software/script/chameleon_cli_main.py @@ -46,7 +46,8 @@ class ChameleonCLI: """ def __init__(self): - self.session = prompt_toolkit.PromptSession() + self.completer = chameleon_utils.CustomNestedCompleter.from_nested_dict(chameleon_cli_unit.root_commands) + self.session = prompt_toolkit.PromptSession(completer=self.completer) # new a device communication instance(only communication) self.device_com = chameleon_com.ChameleonCom() diff --git a/software/script/chameleon_utils.py b/software/script/chameleon_utils.py index 99469b6..20dc0b1 100644 --- a/software/script/chameleon_utils.py +++ b/software/script/chameleon_utils.py @@ -1,5 +1,7 @@ from functools import wraps from typing import Union +from prompt_toolkit.completion import Completer, NestedCompleter, WordCompleter +from prompt_toolkit.document import Document import chameleon_status @@ -84,3 +86,82 @@ class CLITree: CLITree(name=name, fullname=f'{self.fullname} {name}', helptext=helptext, cls=cls)) return cls return decorator + + +class CustomNestedCompleter(NestedCompleter): + """ + Copy of the NestedCompleter class that accepts a CLITree object and + supports meta_dict for descriptions + """ + + def __init__( + self, options, ignore_case: bool = True, meta_dict: dict = {} + ) -> None: + self.options = options + self.ignore_case = ignore_case + self.meta_dict = meta_dict + + def __repr__(self) -> str: + return f"CustomNestedCompleter({self.options!r}, ignore_case={self.ignore_case!r})" + + @classmethod + def from_nested_dict(cls, data): + options = {} + meta_dict = {} + for key, value in data.items(): + if isinstance(value, Completer): + options[key] = value + elif isinstance(value, dict): + options[key] = cls.from_nested_dict(value) + elif isinstance(value, set): + options[key] = cls.from_nested_dict({item: None for item in value}) + elif isinstance(value, CLITree): + options[key] = cls.from_clitree(value) + meta_dict[key] = value.helptext + else: + assert value is None + options[key] = None + + return cls(options, meta_dict=meta_dict) + + @classmethod + def from_clitree(cls, node): + options = {} + meta_dict = {} + + for child_node in node.children: + options[child_node.name] = cls.from_clitree(child_node) + meta_dict[child_node.name] = child_node.helptext + + return cls(options, meta_dict=meta_dict) + + def get_completions(self, document, complete_event): + # Split document. + text = document.text_before_cursor.lstrip() + stripped_len = len(document.text_before_cursor) - len(text) + + # If there is a space, check for the first term, and use a + # subcompleter. + if " " in text: + first_term = text.split()[0] + completer = self.options.get(first_term) + + # If we have a sub completer, use this for the completions. + if completer is not None: + remaining_text = text[len(first_term) :].lstrip() + move_cursor = len(text) - len(remaining_text) + stripped_len + + new_document = Document( + remaining_text, + cursor_position=document.cursor_position - move_cursor, + ) + + yield from completer.get_completions(new_document, complete_event) + + # No space in the input: behave exactly like `WordCompleter`. + else: + completer = WordCompleter( + list(self.options.keys()), ignore_case=self.ignore_case, meta_dict=self.meta_dict + ) + yield from completer.get_completions(document, complete_event) +