Add command autocompletion

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()

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)
+