You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

398 lines
13 KiB

3 years ago
# Written by Sebastian Lohff <>
# Licensed under Apache License 2.0
from prompt_toolkit.application import Application
from prompt_toolkit.buffer import Buffer
from prompt_toolkit.document import Document
from prompt_toolkit.filters import is_searching
3 years ago
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.layout import Layout, Window, HSplit
from prompt_toolkit.layout.controls import BufferControl
3 years ago
from prompt_toolkit.layout.processors import Processor, Transformation
from prompt_toolkit import search
3 years ago
from prompt_toolkit.widgets import SearchToolbar
3 years ago
class _CliMenuHeader:
3 years ago
"""Hold a menu header"""
def __init__(self, text, indent=False):
self.text = text
self.indent = indent
self.focusable = False
class _CliMenuOption:
3 years ago
"""Hold a menu option"""
def __init__(self, text, num, item=None):
3 years ago
self.text = text
self.num = num
self.item = item
3 years ago
self.focusable = True
class CliMenuCursor:
"""Collection of cursors pointing at the active menu item"""
ASCII_CAT = '=^.^='
3 years ago
CAT = '😸'
ARROW = ''
class CliMenuStyle:
"""Style for a menu
Allows to select header, option and selected option color
def __init__(self, option_style='', highlight_style='', header_style=''):
self.option_style = option_style
self.highlight_style = highlight_style
self.header_style = header_style
class CliSelectionStyle:
SQUARE_BRACKETS = ('[x]', '[ ]')
ROUND_BRACKETS = ('(x)', '( )')
CHECKMARK = ('', '')
THUMBS = ('👍', '👎')
SMILEY = ('🙂', '🙁')
SMILEY_EXTREME = ('😁', '😨')
3 years ago
class CliMenuTheme:
BASIC = CliMenuStyle()
BASIC_BOLD = CliMenuStyle(header_style='bold', highlight_style='bold')
RED = CliMenuStyle('#aa0000', '#ee0000', '#aa0000')
CYAN = CliMenuStyle('cyan', 'lightcyan', 'cyan')
BLUE = CliMenuStyle('ansiblue', 'ansired', 'ansiblue')
ANSI_CYAN = CliMenuStyle('ansicyan', 'ansibrightcyan', 'ansicyan')
BOLD_HIGHLIGHT = CliMenuStyle(header_style='bold', highlight_style='bold fg:black bg:white')
3 years ago
class CliMenu:
default_style = CliMenuTheme.BASIC
3 years ago
default_cursor = CliMenuCursor.TRIANGLE
def set_default_style(cls, style):
cls.default_style = style
def set_default_cursor(cls, cursor):
cls.default_cursor = cursor
3 years ago
def __init__(self, options=None, header=None, cursor=None, style=None,
indent=2, dedent_selection=False, initial_pos=0):
3 years ago
self._items = []
self._item_num = 0
self._ran = False
self._success = None
self._pos = 0
self._initial_pos = initial_pos
3 years ago
self._option_indent = indent
self._header_indent = indent
self._dedent_selection = dedent_selection
self._cursor = cursor if cursor is not None else self.default_cursor
self._style = style if style is not None else self.default_style
3 years ago
if header:
self.add_text(header, indent=False)
3 years ago
if options:
for option in options:
if isinstance(option, tuple):
elif isinstance(option, dict):
elif isinstance(option, str):
self.add_option(option, option)
raise ValueError("Option needs to be either tuple, dict or string, found '{}' of type {}"
.format(option, type(option)))
3 years ago
def add_header(self, *args, **kwargs):
return self.add_text(*args, **kwargs)
def add_text(self, title, indent=True):
3 years ago
for text in title.split('\n'):
self._items.append(_CliMenuHeader(text, indent=indent))
3 years ago
def add_option(self, text, item=None):
self._items.append(_CliMenuOption(text, self._item_num, item=item))
3 years ago
self._item_num += 1
def success(self):
3 years ago
if not self._ran:
return self._success
def get_options(self):
return [_item for _item in self._items if isinstance(_item, _CliMenuOption)]
def num_options(self):
return self._item_num
def get_selection(self):
if self.success:
item = self._items[self._pos]
return (item.num, item.item)
return (None, None)
def get_selection_num(self):
return self.get_selection()[0]
def get_selection_item(self):
return self.get_selection()[1]
3 years ago
def _transform_prefix(self, item, lineno, prefix):
return prefix
3 years ago
def _transform_line(self, ti):
if len(list(ti.fragments)) == 0:
return Transformation(ti.fragments)
3 years ago
style, text = list(ti.fragments)[0]
item = self._items[ti.lineno]
s = self._style
# cursor
indent = ''
prefix = ''
if item.focusable:
indent += ' ' * self._option_indent
if ti.lineno == self._pos:
prefix += '{} '.format(self._cursor)
style = s.highlight_style
prefix += ' ' * (len(self._cursor) + 1 + 1 * self._dedent_selection)
style = s.option_style
if item.indent:
indent += ' ' * (self._header_indent + len(self._cursor) + 1)
3 years ago
style = s.header_style
items = [(s if s else style, t) for s, t in ti.fragments]
prefix = self._transform_prefix(item, ti.lineno, prefix)
return Transformation([('', indent), (style, prefix)] + items)
3 years ago
def next_item(self, direction):
if not any(item.focusable for item in self._items):
raise RuntimeError("No focusable item found")
while True:
self._pos = (self._pos + direction) % len(self._items)
# move cursor of buffer along with the selected option
self._buf.cursor_position = self._doc.translate_row_col_to_index(self._pos, 0)
3 years ago
if self._items[self._pos].focusable:
def sync_cursor_to_line(self, line, sync_dir=1):
"""Sync cursor to next fousable item starting on `line`"""
assert sync_dir in (1, -1)
self._pos = line
while not self._items[self._pos].focusable:
self._pos = (self._pos + sync_dir) % len(self._items)
self._buf.cursor_position = self._doc.translate_row_col_to_index(self._pos, 0)
def _get_search_result_lines(self):
"""Get a list of all lines that have a match with the current search result"""
if not self._bufctrl.search_state.text:
return []
idx_list = []
i = 1
while True:
next_idx = self._buf.get_search_position(self._bufctrl.search_state, count=i,
if next_idx in idx_list:
i += 1
lines = []
for idx in idx_list:
line, _ = self._doc.translate_index_to_position(idx)
if line not in lines:
return lines
def _register_extra_kb_cbs(self, kb):
def _preflight(self):
if self._initial_pos < 0 or self._initial_pos >= self._item_num:
raise ValueError("Initial position {} is out of range, needs to be in range of [0, {})"
.format(self._initial_pos, self._item_num))
def _accept(self, event):
self._success = True
def _run(self):
3 years ago
class MenuColorizer(Processor):
def apply_transformation(_self, ti):
return self._transform_line(ti)
# keybindings
self._kb = KeyBindings()
@self._kb.add('q', filter=~is_searching)
3 years ago
def quit(event):
@self._kb.add('down', filter=~is_searching)
@self._kb.add('j', filter=~is_searching)
3 years ago
def down(event):
@self._kb.add('up', filter=~is_searching)
@self._kb.add('k', filter=~is_searching)
3 years ago
def up(event):
@self._kb.add('N', filter=~is_searching)
@self._kb.add('n', filter=~is_searching)
def search_inc(event):
if not self._bufctrl.search_state.text:
search_dir = 1 if == 'n' else -1
sr_lines = self._get_search_result_lines()
if sr_lines:
line = sr_lines[search_dir] if len(sr_lines) > 1 else sr_lines[0]
self.sync_cursor_to_line(line, search_dir)
@self._kb.add('c-m', filter=~is_searching)
@self._kb.add('right', filter=~is_searching)
3 years ago
def accept(event):
3 years ago
@self._kb.add('c-m', filter=is_searching)
def accept_search(event):
new_line, _ = self._doc.translate_index_to_position(self._buf.cursor_position)
self._searchbar = SearchToolbar(ignore_case=True)
3 years ago
text = '\n'.join(map(lambda _x: _x.text, self._items))
self._doc = Document(text, cursor_position=self._pos)
self._buf = Buffer(read_only=True, document=self._doc)
self._bufctrl = BufferControl(self._buf,
3 years ago
split = HSplit([Window(self._bufctrl,
# set initial pos
for _ in range(self._initial_pos + 1):
3 years ago
app = Application(layout=Layout(split),
self._ran = True
class CliMultiMenu(CliMenu):
default_selection_icons = CliSelectionStyle.SQUARE_BRACKETS
def set_default_selector_icons(cls, selection_icons):
cls.default_selection_icons = selection_icons
def __init__(self, *args, selection_icons=None, min_selection_count=0, **kwargs):
self._multi_selected = []
self._min_selection_count = min_selection_count
self._selection_icons = selection_icons if selection_icons is not None else self.default_selection_icons
super().__init__(*args, **kwargs)
def add_option(self, text, item=None, selected=False):
super().add_option(text, item)
if selected:
self._multi_selected.append(len(self._items) - 1)
def get_selection(self):
if self.success:
return [(self._items[n].num, self._items[n].item) for n in self._multi_selected]
return None
def get_selection_num(self):
if self.success:
return [self._items[n].num for n in self._multi_selected]
return None
def get_selection_item(self):
if self.success:
return [self._items[n].item for n in self._multi_selected]
return None
def _register_extra_kb_cbs(self, kb):
@kb.add('space', filter=~is_searching)
@kb.add('right', filter=~is_searching)
def mark(event):
if self._pos not in self._multi_selected:
def _transform_prefix(self, item, lineno, prefix):
if item.focusable:
if lineno in self._multi_selected:
icon = self._selection_icons[0]
icon = self._selection_icons[1]
return "{}{} ".format(prefix, icon)
return prefix
def _preflight(self):
if self._min_selection_count > self._item_num:
raise ValueError("A minimum of {} items was requested for successful selection but only {} exist"
.format(self._min_selection_count, self._item_num))
def _accept(self, event):
if len(self._multi_selected) >= self._min_selection_count:
def cli_select_item(options, header=None, abort_exc=ValueError, abort_text="Selection aborted.", style=None,
"""Helper function to quickly get a selection with just a few arguments"""
menu = CliMenu(header=header, options=options, style=style)
if return_single and menu.num_options == 1:
item = menu.get_options()[0]
return item.num, item.item
if not menu.success:
raise abort_exc(abort_text)
return menu.get_selection()