mirror of
https://github.com/RfidResearchGroup/ChameleonUltra.git
synced 2026-03-31 14:45:49 +00:00
280 lines
9.4 KiB
Python
280 lines
9.4 KiB
Python
import argparse
|
|
from functools import wraps
|
|
from typing import Union
|
|
from prompt_toolkit.completion import Completer, NestedCompleter, WordCompleter
|
|
from prompt_toolkit.completion.base import Completion
|
|
from prompt_toolkit.document import Document
|
|
|
|
import chameleon_status
|
|
|
|
|
|
class ArgsParserError(Exception):
|
|
pass
|
|
|
|
|
|
class ParserExitIntercept(Exception):
|
|
pass
|
|
|
|
|
|
class UnexpectedResponseError(Exception):
|
|
"""
|
|
Unexpected response exception
|
|
"""
|
|
|
|
|
|
class ArgumentParserNoExit(argparse.ArgumentParser):
|
|
"""
|
|
If arg ArgumentParser parse error, we can't exit process,
|
|
we must raise exception to stop parse
|
|
"""
|
|
|
|
def __init__(self, **args):
|
|
super().__init__(*args)
|
|
self.add_help = False
|
|
self.description = "Please enter correct parameters"
|
|
|
|
def exit(self, status: int = ..., message: str or None = ...):
|
|
if message:
|
|
raise ParserExitIntercept(message)
|
|
|
|
def error(self, message: str):
|
|
args = {'prog': self.prog, 'message': message}
|
|
raise ArgsParserError('%(prog)s: error: %(message)s\n' % args)
|
|
|
|
|
|
def expect_response(accepted_responses: Union[int, list[int]]):
|
|
"""
|
|
Decorator for wrapping a Chameleon CMD function to check its response
|
|
for expected return codes and throwing an exception otherwise
|
|
"""
|
|
if isinstance(accepted_responses, int):
|
|
accepted_responses = [accepted_responses]
|
|
|
|
def decorator(func):
|
|
@wraps(func)
|
|
def error_throwing_func(*args, **kwargs):
|
|
ret = func(*args, **kwargs)
|
|
|
|
if ret.status not in accepted_responses:
|
|
if ret.status in chameleon_status.Device and ret.status in chameleon_status.message:
|
|
raise UnexpectedResponseError(
|
|
chameleon_status.message[ret.status])
|
|
else:
|
|
raise UnexpectedResponseError(
|
|
f"Unexpected response and unknown status {ret.status}")
|
|
|
|
return ret
|
|
|
|
return error_throwing_func
|
|
|
|
return decorator
|
|
|
|
|
|
class CLITree:
|
|
"""
|
|
Class holding a
|
|
|
|
:param name: Name of the command (e.g. "set")
|
|
:param help_text: Hint displayed for the command
|
|
:param fullname: Full name of the command that includes previous commands (e.g. "hw mode set")
|
|
:param cls: A BaseCLIUnit instance handling the command
|
|
"""
|
|
|
|
def __init__(self, name=None, help_text=None, fullname=None, children=None, cls=None) -> None:
|
|
self.name: str = name
|
|
self.help_text: str = help_text
|
|
self.fullname: str = fullname if fullname else name
|
|
self.children: list[CLITree] = children if children else list()
|
|
self.cls = cls
|
|
|
|
def subgroup(self, name, help_text=None):
|
|
"""
|
|
Create a child command group
|
|
|
|
:param name: Name of the command group
|
|
:param help_text: Hint displayed for the group
|
|
"""
|
|
child = CLITree(
|
|
name=name, fullname=f'{self.fullname} {name}', help_text=help_text)
|
|
self.children.append(child)
|
|
return child
|
|
|
|
def command(self, name, help_text=None):
|
|
"""
|
|
Create a child command
|
|
|
|
:param name: Name of the command
|
|
:param help_text: Hint displayed for the command
|
|
"""
|
|
def decorator(cls):
|
|
self.children.append(
|
|
CLITree(name=name, fullname=f'{self.fullname} {name}', help_text=help_text, 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):
|
|
if value.cls:
|
|
# CLITree is a standalone command
|
|
options[key] = ArgparseCompleter(value.cls().args_parser())
|
|
else:
|
|
# CLITree is a command group
|
|
options[key] = cls.from_clitree(value)
|
|
meta_dict[key] = value.help_text
|
|
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:
|
|
if child_node.cls and child_node.cls().args_parser():
|
|
# CLITree is a standalone command with arguments
|
|
options[child_node.name] = ArgparseCompleter(
|
|
child_node.cls().args_parser())
|
|
else:
|
|
# CLITree is a command group
|
|
options[child_node.name] = cls.from_clitree(child_node)
|
|
meta_dict[child_node.name] = child_node.help_text
|
|
|
|
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 sub_completer.
|
|
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)
|
|
|
|
|
|
class ArgparseCompleter(Completer):
|
|
"""
|
|
Completer instance for autocompletion of ArgumentParser arguments
|
|
|
|
:param parser: ArgumentParser instance
|
|
"""
|
|
|
|
def __init__(self, parser) -> None:
|
|
self.parser: ArgumentParserNoExit = parser
|
|
|
|
def check_tokens(self, parsed, unparsed):
|
|
suggestions = {}
|
|
|
|
def check_arg(tokens):
|
|
return tokens and tokens[0].startswith('-')
|
|
|
|
if not parsed and not unparsed:
|
|
# No tokens detected, just show all flags
|
|
for action in self.parser._actions:
|
|
for opt in action.option_strings:
|
|
suggestions[opt] = action.help
|
|
return [], [], suggestions
|
|
|
|
token = unparsed.pop(0)
|
|
|
|
for action in self.parser._actions:
|
|
if any(opt == token for opt in action.option_strings):
|
|
# Argument fully matches the token
|
|
parsed.append(token)
|
|
|
|
if action.choices:
|
|
# Autocomplete with choices
|
|
if unparsed:
|
|
# Autocomplete values
|
|
value = unparsed.pop(0)
|
|
for choice in action.choices:
|
|
if str(choice).startswith(value):
|
|
suggestions[str(choice)] = None
|
|
|
|
parsed.append(value)
|
|
|
|
if check_arg(unparsed):
|
|
parsed, unparsed, suggestions = self.check_tokens(
|
|
parsed, unparsed)
|
|
|
|
else:
|
|
# Show all possible values
|
|
for choice in action.choices:
|
|
suggestions[str(choice)] = None
|
|
|
|
break
|
|
else:
|
|
# No choices, process further arguments
|
|
if check_arg(unparsed):
|
|
parsed, unparsed, suggestions = self.check_tokens(
|
|
parsed, unparsed)
|
|
break
|
|
elif any(opt.startswith(token) for opt in action.option_strings):
|
|
for opt in action.option_strings:
|
|
if opt.startswith(token):
|
|
suggestions[opt] = action.help
|
|
|
|
if suggestions:
|
|
unparsed.insert(0, token)
|
|
|
|
return parsed, unparsed, suggestions
|
|
|
|
def get_completions(self, document, complete_event):
|
|
text = document.text_before_cursor
|
|
word_before_cursor = document.text_before_cursor.split(' ')[-1]
|
|
|
|
_, _, suggestions = self.check_tokens(list(), text.split())
|
|
|
|
for key, suggestion in suggestions.items():
|
|
yield Completion(key, -len(word_before_cursor), display=key, display_meta=suggestion)
|