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)