1""" 
    2Nestedcompleter for completion of hierarchical data structures. 
    3""" 
    4 
    5from __future__ import annotations 
    6 
    7from typing import Any, Iterable, Mapping, Set, Union 
    8 
    9from prompt_toolkit.completion import CompleteEvent, Completer, Completion 
    10from prompt_toolkit.completion.word_completer import WordCompleter 
    11from prompt_toolkit.document import Document 
    12 
    13__all__ = ["NestedCompleter"] 
    14 
    15# NestedDict = Mapping[str, Union['NestedDict', Set[str], None, Completer]] 
    16NestedDict = Mapping[str, Union[Any, Set[str], None, Completer]] 
    17 
    18 
    19class NestedCompleter(Completer): 
    20    """ 
    21    Completer which wraps around several other completers, and calls any the 
    22    one that corresponds with the first word of the input. 
    23 
    24    By combining multiple `NestedCompleter` instances, we can achieve multiple 
    25    hierarchical levels of autocompletion. This is useful when `WordCompleter` 
    26    is not sufficient. 
    27 
    28    If you need multiple levels, check out the `from_nested_dict` classmethod. 
    29    """ 
    30 
    31    def __init__( 
    32        self, options: dict[str, Completer | None], ignore_case: bool = True 
    33    ) -> None: 
    34        self.options = options 
    35        self.ignore_case = ignore_case 
    36 
    37    def __repr__(self) -> str: 
    38        return f"NestedCompleter({self.options!r}, ignore_case={self.ignore_case!r})" 
    39 
    40    @classmethod 
    41    def from_nested_dict(cls, data: NestedDict) -> NestedCompleter: 
    42        """ 
    43        Create a `NestedCompleter`, starting from a nested dictionary data 
    44        structure, like this: 
    45 
    46        .. code:: 
    47 
    48            data = { 
    49                'show': { 
    50                    'version': None, 
    51                    'interfaces': None, 
    52                    'clock': None, 
    53                    'ip': {'interface': {'brief'}} 
    54                }, 
    55                'exit': None 
    56                'enable': None 
    57            } 
    58 
    59        The value should be `None` if there is no further completion at some 
    60        point. If all values in the dictionary are None, it is also possible to 
    61        use a set instead. 
    62 
    63        Values in this data structure can be a completers as well. 
    64        """ 
    65        options: dict[str, Completer | None] = {} 
    66        for key, value in data.items(): 
    67            if isinstance(value, Completer): 
    68                options[key] = value 
    69            elif isinstance(value, dict): 
    70                options[key] = cls.from_nested_dict(value) 
    71            elif isinstance(value, set): 
    72                options[key] = cls.from_nested_dict(dict.fromkeys(value)) 
    73            else: 
    74                assert value is None 
    75                options[key] = None 
    76 
    77        return cls(options) 
    78 
    79    def get_completions( 
    80        self, document: Document, complete_event: CompleteEvent 
    81    ) -> Iterable[Completion]: 
    82        # Split document. 
    83        text = document.text_before_cursor.lstrip() 
    84        stripped_len = len(document.text_before_cursor) - len(text) 
    85 
    86        # If there is a space, check for the first term, and use a 
    87        # subcompleter. 
    88        if " " in text: 
    89            first_term = text.split()[0] 
    90            completer = self.options.get(first_term) 
    91 
    92            # If we have a sub completer, use this for the completions. 
    93            if completer is not None: 
    94                remaining_text = text[len(first_term) :].lstrip() 
    95                move_cursor = len(text) - len(remaining_text) + stripped_len 
    96 
    97                new_document = Document( 
    98                    remaining_text, 
    99                    cursor_position=document.cursor_position - move_cursor, 
    100                ) 
    101 
    102                yield from completer.get_completions(new_document, complete_event) 
    103 
    104        # No space in the input: behave exactly like `WordCompleter`. 
    105        else: 
    106            completer = WordCompleter( 
    107                list(self.options.keys()), ignore_case=self.ignore_case 
    108            ) 
    109            yield from completer.get_completions(document, complete_event)