From 66a13dcd29971ca7701acd02eaf8d4b8a65e9871 Mon Sep 17 00:00:00 2001 From: UnknownUser95 Date: Sun, 7 Jul 2024 03:46:54 +0200 Subject: [PATCH] refactor selection methods --- src/pickpack/__init__.py | 155 ++++++++++++++++++---------------- src/pickpack/anytree_utils.py | 5 +- 2 files changed, 84 insertions(+), 76 deletions(-) diff --git a/src/pickpack/__init__.py b/src/pickpack/__init__.py index a8093a5..7697243 100644 --- a/src/pickpack/__init__.py +++ b/src/pickpack/__init__.py @@ -20,8 +20,8 @@ KEYS_DOWN = (curses.KEY_DOWN, ord('j')) KEYS_SELECT = (curses.KEY_RIGHT, ord(' ')) -NodeWithIndex: TypeAlias = tuple[str | Node, int] | list[tuple[str | Node, int]] -NodeNameOnly: TypeAlias = str | Node | list[Node | str] +NodeWithIndex: TypeAlias = tuple[str | Node, int] +NodeNameOnly: TypeAlias = str | Node AnyNode: TypeAlias = NodeNameOnly | NodeWithIndex _T = TypeVar("_T") @@ -80,6 +80,10 @@ class PickPacker: custom_handlers: dict[int, Callable[[PickPacker], Any]] = field(init=False, default_factory=dict) index: int = field(init=False, default=0) scroll_top: int = field(init=False, default=0) + + @property + def name_only(self) -> bool: + return self.output_format == OutputMode.nameonly or self.output_format == OutputMode.nameindex def __post_init__(self): # Check for correct number of elements @@ -192,94 +196,99 @@ def mark_index(self) -> None: else: self.all_selected.append(self.index) self.add_relatives_index() - - def get_selected_noindex(self) -> NodeNameOnly: - nameonly = self.output_format == OutputMode.nameonly - if self.multiselect: - return_tuples: list[Node] = [] + + def no_index_multi_select(self) -> list[NodeNameOnly]: + nodes: list[NodeNameOnly] = [] + + for selected in self.all_selected: + nodes.append(find_by_index(self.options.node, selected)) + + if self.output_leaves_only: + nodes = [n for n in nodes if n.children] + if self.name_only: + nodes = [n.name for n in nodes] + + return nodes + + def no_index_single_select(self) -> NodeNameOnly | list[NodeNameOnly]: + node: Node = find_by_index(self.options.node, self.index) + + if self.output_leaves_only or self.singleselect_output_include_children: if self.output_leaves_only: - for selected in self.all_selected: - node = find_by_index(self.options.node, selected) - return_tuples.extend([leaf for leaf in get_leaves_only(node) if leaf not in return_tuples]) + nodes: list[str | Node] = get_leaves_only(node) else: - for selected in self.all_selected: - if nameonly: - return_tuples.append(find_by_index(self.options.node, selected).name) - else: - return_tuples.append(find_by_index(self.options.node, selected)) - return return_tuples + nodes: list[str | Node] = get_descendants(node) + + if self.name_only: + nodes = [n.name for n in nodes] + + return nodes + + if self.name_only: + node = node.name + + return node + + def get_selected_no_index(self) -> NodeNameOnly | list[NodeNameOnly]: + if self.multiselect: + return self.no_index_multi_select() else: - node = find_by_index(self.options.node, self.index) + return self.no_index_single_select() + + def index_multi_select(self) -> list[NodeWithIndex]: + nodes: list[tuple[str | Node, int]] = [] + + if self.output_leaves_only: + for selected in self.all_selected: + node: Node = find_by_index(self.options.node, selected) + nodes.extend([(leaf, leaf.index) for leaf in get_leaves_only(node) if (leaf, leaf.index) not in nodes]) + else: + for selected in self.all_selected: + nodes.append((find_by_index(self.options.node, selected), selected)) + + if self.name_only: + nodes = [(n[0].name, n[1]) for n in nodes] + + return nodes + + def index_single_select(self) -> NodeWithIndex | list[NodeWithIndex]: + node: Node = find_by_index(self.options.node, self.index) + + if self.output_leaves_only or self.singleselect_output_include_children: if self.output_leaves_only: - if nameonly: - return [leaf.name for leaf in get_leaves_only(node)] - else: - return get_leaves_only(node) + nodes: list[Node] = get_leaves_only(node) else: - if self.singleselect_output_include_children: - descendants = get_descendants(node) - if nameonly: - return [node.name for node in descendants] - else: - return descendants - else: - if nameonly: - return node.name - else: - return node + nodes: list[Node] = get_descendants(node) + + index_nodes: list[tuple[str | Node, int]] = [(n, n.index) for n in nodes] + + if self.name_only: + index_nodes = [(n[0].name, n[1]) for n in index_nodes] + + return index_nodes + + if self.name_only: + node = node.name + + return node, self.index - def get_selected_withindex(self) -> NodeWithIndex: + def get_selected_with_index(self) -> NodeWithIndex | list[NodeWithIndex]: """return the current selected option as a tuple: (option, index) - or as a list of tuples (in case multiselect==True) + or as a list of tuples (in case multiselect is True) """ - nameonly = bool(self.output_format) - if self.multiselect: - return_tuples: list[tuple[str | Node, int]] = [] - if self.output_leaves_only: - for selected in self.all_selected: - node = find_by_index(self.options.node, selected) - if nameonly: - return_tuples.extend([(leaf.name, leaf.index) for leaf in get_leaves_only(node) if (leaf.name, leaf.index) not in return_tuples]) - else: - return_tuples.extend([(leaf, leaf.index) for leaf in get_leaves_only(node) if (leaf, leaf.index) not in return_tuples]) - - else: - if nameonly: - for selected in self.all_selected: - return_tuples.append((find_by_index(self.options.node, selected).name, selected)) - else: - for selected in self.all_selected: - return_tuples.append((find_by_index(self.options.node, selected), selected)) - return return_tuples + return self.index_multi_select() else: - node = find_by_index(self.options.node, self.index) - if self.singleselect_output_include_children: - if self.output_leaves_only: - if nameonly: - return [(leaf.name, leaf.index) for leaf in get_leaves_only(node)] - else: - return [(leaf, leaf.index) for leaf in get_leaves_only(node)] - else: - descendants = get_descendants(node) - if nameonly: - return [(node.name, node.index) for node in descendants] - else: - return [(node, node.index) for node in descendants] - else: - if nameonly: - return node.name, self.index - else: - return node, self.index + return self.index_single_select() def get_selected(self) -> AnyNode: """return the current selected option as a tuple: (option, index) or as a list of tuples (in case multiselect is True) """ if self.output_format == OutputMode.nodeindex or self.output_format == OutputMode.nameindex: - return self.get_selected_withindex() + return self.get_selected_with_index() elif self.output_format == OutputMode.nameonly or self.output_format == OutputMode.nodeonly: - return self.get_selected_noindex() + return self.get_selected_no_index() def get_title_lines(self) -> list[str]: if self.title: diff --git a/src/pickpack/anytree_utils.py b/src/pickpack/anytree_utils.py index 0e35f10..08d601e 100644 --- a/src/pickpack/anytree_utils.py +++ b/src/pickpack/anytree_utils.py @@ -54,7 +54,6 @@ def get_leaves_only(node: Node, leaves: list[Node] = None) -> list[Node]: if node.children: for child in node.children: leaves = get_leaves_only(child, leaves) - else: - if node not in leaves: - leaves.append(node) + elif node not in leaves: + leaves.append(node) return leaves