Files
AirlockTools/widgets/policytreewidget.py
T

207 lines
7.4 KiB
Python

import logging
from rich.text import Text
from textual.containers import Horizontal, Vertical
from textual.widget import Widget
from textual.widgets import Input, OptionList, Static, Tree
from textual.widgets.option_list import Option
logger = logging.getLogger(__name__)
class PolicyTreeWidget(Widget):
"""Widget for displaying and searching a hierarchical policy tree."""
def __init__(self, policies, devices):
super().__init__()
self.policies = policies
self.devices = devices
self.last_highlighted_node = None
def compose(self):
# Left: Policy Tree
policy_tree = Tree("Policies", id="policy_tree")
policy_tree.styles.width = "2fr"
policy_tree.styles.height = "100%"
# Right: Search + Details
label = Static("Device Search:")
search_box = Input(
placeholder="Search policies or devices...", id="tree_search"
)
details_pane = Static("", id="details_pane")
with Horizontal():
yield policy_tree
with Vertical() as right_pane:
right_pane.styles.width = "3fr"
yield label
yield search_box
yield details_pane
def on_mount(self) -> None:
"""Build the tree after mounting."""
self._build_tree()
def _build_tree(self) -> None:
"""Build the policy tree structure."""
policy_tree = self.query_one("#policy_tree", Tree)
node_map = {}
# Top-level policies
for policy in self.policies:
if policy.parent == "global-policy-settings":
node = policy_tree.root.add(label=policy.name, data=policy)
node_map[policy.groupid] = node
# Child policies
for policy in self.policies:
parent_id = policy.parent
if parent_id in node_map:
parent_node = node_map[parent_id]
node = parent_node.add(label=policy.name, data=policy)
node_map[policy.groupid] = node
# Devices under policies
for device in self.devices:
group_id = device.groupid
if group_id in node_map:
parent_node = node_map[group_id]
label = device.hostname
parent_node.add(label=label, data=device)
def _collect_tree_nodes(self, node, all_nodes):
"""Helper to recursively collect all nodes from a tree."""
all_nodes.append(node)
for child in node.children:
self._collect_tree_nodes(child, all_nodes)
def _remove_match_selector(self):
"""Safely remove match selector widgets."""
try:
existing = self.query("#match_selector")
for widget in existing:
if widget.is_attached:
widget.remove()
except Exception as exc:
logger.debug("Failed to remove match_selector: %s", exc)
def on_tree_node_selected(self, message: Tree.NodeSelected) -> None:
"""Handle tree node selection."""
node = message.node
data = node.data
details_pane = self.query_one("#details_pane", Static)
# Reset previous highlight
if self.last_highlighted_node is not None:
original_label = str(self.last_highlighted_node.label).strip()
# Remove any styling
if isinstance(self.last_highlighted_node.label, Text):
original_label = self.last_highlighted_node.label.plain
self.last_highlighted_node.set_label(original_label)
# Apply highlight to current node
label_text = str(node.label).strip()
if isinstance(node.label, Text):
label_text = node.label.plain
highlighted_label = Text(label_text, style="reverse bold")
node.set_label(highlighted_label)
self.last_highlighted_node = node
# Update details pane
if data:
# Work with dataclass objects using __dict__
details = "\n".join(
f"{key}: {value}" for key, value in data.__dict__.items()
)
else:
details = f"Selected: {node.label}"
details_pane.update(details)
# Stop event from bubbling
message.stop()
def on_input_submitted(self, message: Input.Submitted) -> None:
"""Handle search input submission."""
# Remove existing match selector FIRST
self._remove_match_selector()
query = message.value.strip().lower()
tree = self.query_one("#policy_tree", Tree)
details_pane = self.query_one("#details_pane", Static)
all_nodes = []
self._collect_tree_nodes(tree.root, all_nodes)
label_to_node = {}
for node in all_nodes:
label_text = str(node.label).lower()
label_to_node[label_text] = node
if node.data:
# Use __dict__ for dataclass objects
data_dict = (
node.data.__dict__ if hasattr(node.data, "__dict__") else node.data
)
for key, value in data_dict.items():
if isinstance(value, str):
label_to_node[value.lower()] = node
# Wildcard-style substring match
matches = sorted([label for label in label_to_node if query in label])
if matches:
# Try to reuse existing match_selector or create new one
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.clear_options()
option_list.display = True # Ensure it's visible
except:
option_list = OptionList(id="match_selector")
# Mount to the details pane's parent (the Vertical container)
details_pane.parent.mount(option_list)
for label in matches:
option_list.add_option(Option(label, id=f"match_{label}"))
details_pane.update(f"Found {len(matches)} matches. Select one below.")
else:
# Hide or remove the match_selector when no matches
self._remove_match_selector()
details_pane.update("No matches found.")
def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None:
"""Handle selection from search results."""
selected_id = event.option.id.replace("match_", "")
tree = self.query_one("#policy_tree", Tree)
details_pane = self.query_one("#details_pane", Static)
# Find the node
all_nodes = []
self._collect_tree_nodes(tree.root, all_nodes)
label_to_node = {str(node.label).lower(): node for node in all_nodes}
match_node = label_to_node.get(selected_id.lower())
if match_node:
# Expand path (original working logic)
node = match_node
path = []
while node:
path.insert(0, node)
node = node.parent
for node in path:
node.expand()
tree.select_node(match_node)
tree.scroll_to_node(match_node)
match_node.set_label(Text(str(match_node.label), style="reverse bold"))
details_pane.update(f"Selected: {match_node.label}")
# Remove the match_selector after selection
try:
option_list = self.query_one("#match_selector", OptionList)
option_list.remove()
except:
pass