"""Nested completer for completion of OpenBB hierarchical data structures.""" from typing import ( Any, Callable, Dict, Iterable, List, Mapping, Optional, Pattern, Set, Union, ) from prompt_toolkit.completion import CompleteEvent, Completer, Completion from prompt_toolkit.document import Document from prompt_toolkit.formatted_text import AnyFormattedText NestedDict = Mapping[str, Union[Any, Set[str], None, Completer]] # pylint: disable=too-many-arguments,global-statement,too-many-branches,global-variable-not-assigned class WordCompleter(Completer): """Simple autocompletion on a list of words. :param words: List of words or callable that returns a list of words. :param ignore_case: If True, case-insensitive completion. :param meta_dict: Optional dict mapping words to their meta-text. (This should map strings to strings or formatted text.) :param WORD: When True, use WORD characters. :param sentence: When True, don't complete by comparing the word before the cursor, but by comparing all the text before the cursor. In this case, the list of words is just a list of strings, where each string can contain spaces. (Can not be used together with the WORD option.) :param match_middle: When True, match not only the start, but also in the middle of the word. :param pattern: Optional compiled regex for finding the word before the cursor to complete. When given, use this regex pattern instead of default one (see document._FIND_WORD_RE) """ def __init__( self, words: Union[List[str], Callable[[], List[str]]], ignore_case: bool = False, display_dict: Optional[Mapping[str, AnyFormattedText]] = None, meta_dict: Optional[Mapping[str, AnyFormattedText]] = None, WORD: bool = True, sentence: bool = False, match_middle: bool = False, pattern: Optional[Pattern[str]] = None, ) -> None: """Initialize the WordCompleter.""" assert not (WORD and sentence) # noqa: S101 self.words = words self.ignore_case = ignore_case self.display_dict = display_dict or {} self.meta_dict = meta_dict or {} self.WORD = WORD self.sentence = sentence self.match_middle = match_middle self.pattern = pattern def get_completions( self, document: Document, _complete_event: CompleteEvent, ) -> Iterable[Completion]: """Get completions.""" # Get list of words. words = self.words if callable(words): words = words() # Get word/text before cursor. if self.sentence: word_before_cursor = document.text_before_cursor else: word_before_cursor = document.get_word_before_cursor( WORD=self.WORD, pattern=self.pattern ) if ( "--" in document.text_before_cursor and document.text_before_cursor.rfind(" --") >= document.text_before_cursor.rfind(" -") ): word_before_cursor = f'--{document.text_before_cursor.split("--")[-1]}' elif f"--{word_before_cursor}" == document.text_before_cursor: word_before_cursor = document.text_before_cursor if self.ignore_case: word_before_cursor = word_before_cursor.lower() def word_matches(word: str) -> bool: """Set True when the word before the cursor matches.""" if self.ignore_case: word = word.lower() if self.match_middle: return word_before_cursor in word return word.startswith(word_before_cursor) for a in words: if word_matches(a): display = self.display_dict.get(a, a) display_meta = self.meta_dict.get(a, "") yield Completion( text=a, start_position=-len(word_before_cursor), display=display, display_meta=display_meta, ) 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. """ complementary: List = list() def __init__( self, options: Dict[str, Optional[Completer]], ignore_case: bool = True ) -> None: """Initialize the NestedCompleter.""" self.flags_processed: List = list() self.original_options = options self.options = options self.ignore_case = ignore_case self.complementary = list() def __repr__(self) -> str: """Return string representation of NestedCompleter.""" return f"NestedCompleter({self.options!r}, ignore_case={self.ignore_case!r})" @classmethod def from_nested_dict(cls, data: dict) -> "NestedCompleter": """Create a `NestedCompleter`. It starts 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, Any] = {} 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(key, str) and isinstance(value, str): options[key] = options[value] else: assert value is None # noqa: S101 options[key] = None for items in cls.complementary: if items[0] in options: options[items[1]] = options[items[0]] elif items[1] in options: options[items[0]] = options[items[1]] return cls(options) def get_completions( # noqa: PLR0912 self, document: Document, complete_event: CompleteEvent ) -> Iterable[Completion]: """Get completions.""" # Split document. cmd = "" text = document.text_before_cursor.lstrip() if " " in text: cmd = text.split(" ")[0] if "-" in text: if text.rfind("--") == -1 or text.rfind("-") - 1 > text.rfind("--"): unprocessed_text = "-" + text.split("-")[-1] else: unprocessed_text = "--" + text.split("--")[-1] else: unprocessed_text = text stripped_len = len(document.text_before_cursor) - len(text) # Check if there are multiple flags for the same command if self.complementary: for same_flags in self.complementary: if ( same_flags[0] in self.flags_processed and same_flags[1] not in self.flags_processed ) or ( same_flags[1] in self.flags_processed and same_flags[0] not in self.flags_processed ): if same_flags[0] in self.flags_processed: self.flags_processed.append(same_flags[1]) elif same_flags[1] in self.flags_processed: self.flags_processed.append(same_flags[0]) if cmd: self.options = { k: self.original_options.get(cmd).options[k] # type: ignore for k in self.original_options.get(cmd).options # type: ignore if k not in self.flags_processed } else: self.options = { k: self.original_options[k] for k in self.original_options if k not in self.flags_processed } # If there is a space, check for the first term, and use a subcompleter. if " " in unprocessed_text: first_term = unprocessed_text.split()[0] # user is updating one of the values if unprocessed_text[-1] != " ": self.flags_processed = [ flag for flag in self.flags_processed if flag != first_term ] if self.complementary: for same_flags in self.complementary: if ( same_flags[0] in self.flags_processed and same_flags[1] not in self.flags_processed ) or ( same_flags[1] in self.flags_processed and same_flags[0] not in self.flags_processed ): if same_flags[0] in self.flags_processed: self.flags_processed.remove(same_flags[0]) elif same_flags[1] in self.flags_processed: self.flags_processed.remove(same_flags[1]) if cmd and self.original_options.get(cmd): self.options = self.original_options else: self.options = { k: self.original_options[k] for k in self.original_options if k not in self.flags_processed } if "-" not in text: completer = self.options.get(first_term) elif cmd in self.options and self.options.get(cmd): completer = self.options.get(cmd).options.get(first_term) # type: ignore else: completer = self.options.get(first_term) # If we have a sub completer, use this for the completions. if completer is not None: remaining_text = unprocessed_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, ) # Provides auto-completion but if user doesn't take it still keep going if " " in new_document.text: if ( new_document.text in [f"{opt} " for opt in self.options] or unprocessed_text[-1] == " " ): self.flags_processed.append(first_term) if cmd: self.options = { k: self.original_options.get(cmd).options[k] # type: ignore for k in self.original_options.get(cmd).options # type: ignore if k not in self.flags_processed } else: self.options = { k: self.original_options[k] for k in self.original_options if k not in self.flags_processed } # In case the users inputs a single boolean flag elif not completer.options: # type: ignore self.flags_processed.append(first_term) if self.complementary: for same_flags in self.complementary: if ( same_flags[0] in self.flags_processed and same_flags[1] not in self.flags_processed ) or ( same_flags[1] in self.flags_processed and same_flags[0] not in self.flags_processed ): if same_flags[0] in self.flags_processed: self.flags_processed.append(same_flags[1]) elif same_flags[1] in self.flags_processed: self.flags_processed.append(same_flags[0]) if cmd: self.options = { k: self.original_options.get(cmd).options[k] # type: ignore for k in self.original_options.get(cmd).options # type: ignore if k not in self.flags_processed } else: self.options = { k: self.original_options[k] for k in self.original_options if k not in self.flags_processed } else: # This is a NestedCompleter yield from completer.get_completions(new_document, complete_event) # No space in the input: behave exactly like `WordCompleter`. else: # check if the prompt has been updated in the meantime if " " in text or "-" in text: actual_flags_processed = [ flag for flag in self.flags_processed if flag in text ] if self.complementary: for same_flags in self.complementary: if ( same_flags[0] in actual_flags_processed and same_flags[1] not in actual_flags_processed ) or ( same_flags[1] in actual_flags_processed and same_flags[0] not in actual_flags_processed ): if same_flags[0] in actual_flags_processed: actual_flags_processed.append(same_flags[1]) elif same_flags[1] in actual_flags_processed: actual_flags_processed.append(same_flags[0]) if len(actual_flags_processed) < len(self.flags_processed): self.flags_processed = actual_flags_processed if cmd: self.options = { k: self.original_options.get(cmd).options[k] # type: ignore for k in self.original_options.get(cmd).options # type: ignore if k not in self.flags_processed } else: self.options = { k: self.original_options[k] for k in self.original_options if k not in self.flags_processed } command = self.options.get(cmd) options = command.options if command else {} # type: ignore command_options = [f"{cmd} {opt}" for opt in options] text_list = [text in val for val in command_options] if cmd and cmd in self.options and text_list: completer = WordCompleter( list(self.options.get(cmd).options.keys()), # type: ignore ignore_case=self.ignore_case, ) elif bool([val for val in self.options if text in val]): completer = WordCompleter( list(self.options.keys()), ignore_case=self.ignore_case ) else: # The user has delete part of the first command and we need to reset options if bool([val for val in self.original_options if text in val]): self.options = self.original_options self.flags_processed = list() completer = WordCompleter( list(self.options.keys()), ignore_case=self.ignore_case ) # This is a WordCompleter yield from completer.get_completions(document, complete_event)