""" Nestedcompleter for completion of hierarchical data structures. """ from typing import Any, Dict, Iterable, Mapping, Optional, Set, Union from prompt_toolkit.completion import CompleteEvent, Completer, Completion from prompt_toolkit.completion.word_completer import WordCompleter from prompt_toolkit.document import Document __all__ = ["NestedCompleter"] # NestedDict = Mapping[str, Union['NestedDict', Set[str], None, Completer]] NestedDict = Mapping[str, Union[Any, Set[str], None, Completer]] class NestedCompleter(Completer): """ Completer which wraps around several other completers, and calls any the one that corresponds with the first word of the input. By combining multiple `NestedCompleter` instances, we can achieve multiple hierarchical levels of autocompletion. This is useful when `WordCompleter` is not sufficient. If you need multiple levels, check out the `from_nested_dict` classmethod. """ def __init__( self, options: Dict[str, Optional[Completer]], ignore_case: bool = True ) -> None: self.options = options self.ignore_case = ignore_case def __repr__(self) -> str: return "NestedCompleter(%r, ignore_case=%r)" % (self.options, self.ignore_case) @classmethod def from_nested_dict(cls, data: NestedDict) -> "NestedCompleter": """ Create a `NestedCompleter`, starting from a nested dictionary data structure, like this: .. code:: data = { 'show': { 'version': None, 'interfaces': None, 'clock': None, 'ip': {'interface': {'brief'}} }, 'exit': None 'enable': None } The value should be `None` if there is no further completion at some point. If all values in the dictionary are None, it is also possible to use a set instead. Values in this data structure can be a completers as well. """ options: Dict[str, Optional[Completer]] = {} 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}) else: assert value is None options[key] = None return cls(options) def get_completions( self, document: Document, complete_event: CompleteEvent ) -> Iterable[Completion]: # 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, ) for c in completer.get_completions(new_document, complete_event): yield c # No space in the input: behave exactly like `WordCompleter`. else: completer = WordCompleter( list(self.options.keys()), ignore_case=self.ignore_case ) for c in completer.get_completions(document, complete_event): yield c