DynamicAbstractionSystem/ui/tree_widget.py

320 lines
12 KiB
Python

"""
Tree widget for displaying hierarchical inspector view.
Handles rendering, interaction, and navigation of tree nodes.
"""
import pygame
import pygame_gui
from pygame_gui.core import UIElement
from typing import List, Optional, Tuple, Any
from ui.inspector_tree import TreeNode, SimulationNode, TreeSelectionManager
class TreeWidget(UIElement):
"""Interactive tree widget for hierarchical entity inspection."""
def __init__(self, relative_rect: pygame.Rect, manager: pygame_gui.UIManager, world: Any):
super().__init__(relative_rect, manager, container=None,
starting_height=1, layer_thickness=1)
self.world = world
self.font = pygame.font.Font(None, 16)
self.small_font = pygame.font.Font(None, 12)
# Tree structure
self.root_node = SimulationNode(world)
self.visible_nodes: List[TreeNode] = []
self.all_nodes: List[TreeNode] = []
# Selection management
self.selection_manager = TreeSelectionManager()
# Visual properties
self.node_height = 20
self.expand_collapse_width = 20
self.icon_size = 8
self.text_color = (200, 200, 200)
self.selected_color = (50, 100, 150)
self.hover_color = (60, 60, 80)
self.expand_icon_color = (150, 150, 150)
# Interaction state
self.hovered_node: Optional[TreeNode] = None
self.drag_start_node: Optional[TreeNode] = None
self.is_dragging = False
# Scrolling
self.scroll_offset = 0
self.total_height = 0
# Initialize tree structure
self._update_tree_structure()
self._update_visible_nodes()
def _update_tree_structure(self) -> None:
"""Update the tree structure based on current world state."""
self.root_node.update_entity_counts()
self.all_nodes = self._build_all_nodes_list(self.root_node)
def _build_all_nodes_list(self, node: TreeNode) -> List[TreeNode]:
"""Build a flat list of all nodes in the tree."""
nodes = [node]
for child in node.children:
nodes.extend(self._build_all_nodes_list(child))
return nodes
def _update_visible_nodes(self) -> None:
"""Update the list of visible nodes based on expanded state."""
self.visible_nodes = [self.root_node]
self.visible_nodes.extend(self.root_node.get_all_visible_descendants())
self._update_node_layout()
self.total_height = len(self.visible_nodes) * self.node_height
def _update_node_layout(self) -> None:
"""Update the layout rectangles for all visible nodes."""
y_offset = 0
for node in self.visible_nodes:
node.rect = pygame.Rect(
0,
y_offset - self.scroll_offset,
self.rect.width,
self.node_height
)
y_offset += self.node_height
def _get_node_at_position_local(self, local_pos: Tuple[int, int]) -> Optional[TreeNode]:
"""Get the node at the given local position (already converted)."""
for node in self.visible_nodes:
if node.rect.collidepoint(local_pos):
return node
return None
def _get_node_at_position(self, pos: Tuple[int, int]) -> Optional[TreeNode]:
"""Get the node at the given screen position."""
# Convert screen coordinates to local tree widget coordinates
local_pos = (pos[0] - self.rect.x, pos[1] - self.rect.y)
return self._get_node_at_position_local(local_pos)
def _get_expand_collapse_rect(self, node: TreeNode) -> pygame.Rect:
"""Get the rectangle for the expand/collapse icon."""
indent = node.get_indent()
return pygame.Rect(
indent + 5,
node.rect.y + (self.node_height - self.icon_size) // 2,
self.icon_size,
self.icon_size
)
def _is_click_on_expand_collapse_local(self, node: TreeNode, local_pos: Tuple[int, int]) -> bool:
"""Check if a click is on the expand/collapse icon (local coordinates)."""
if not node.can_expand():
return False
expand_rect = self._get_expand_collapse_rect(node)
return expand_rect.collidepoint(local_pos)
def _is_click_on_expand_collapse(self, node: TreeNode, pos: Tuple[int, int]) -> bool:
"""Check if a click is on the expand/collapse icon."""
if not node.can_expand():
return False
local_pos = (pos[0] - self.rect.x, pos[1] - self.rect.y)
expand_rect = self._get_expand_collapse_rect(node)
return expand_rect.collidepoint(local_pos)
def _render_node(self, surface: pygame.Surface, node: TreeNode) -> None:
"""Render a single tree node."""
if not node.rect.colliderect(surface.get_rect()):
return
# Background
if node.is_selected:
pygame.draw.rect(surface, self.selected_color, node.rect)
elif node == self.hovered_node:
pygame.draw.rect(surface, self.hover_color, node.rect)
# Expand/collapse icon
if node.can_expand():
icon_rect = self._get_expand_collapse_rect(node)
icon_color = self.expand_icon_color
# Draw + or - icon
center_x = icon_rect.centerx
center_y = icon_rect.centery
size = 3
# Horizontal line
pygame.draw.line(surface, icon_color,
(center_x - size, center_y),
(center_x + size, center_y), 1)
# Vertical line (only for collapsed nodes)
if not node.is_expanded:
pygame.draw.line(surface, icon_color,
(center_x, center_y - size),
(center_x, center_y + size), 1)
# Node text
indent = node.get_indent()
text_x = indent + (self.expand_collapse_width if node.can_expand() else 10)
text_surface = self.font.render(node.get_display_text(), True, self.text_color)
text_rect = pygame.Rect(
text_x,
node.rect.y + (self.node_height - text_surface.get_height()) // 2,
text_surface.get_width(),
text_surface.get_height()
)
# Clip text if it's too wide
max_text_width = self.rect.width - text_x - 10
if text_rect.width > max_text_width:
# Truncate text
text = node.get_display_text()
while text and text_surface.get_width() > max_text_width - 20:
text = text[:-1]
text_surface = self.font.render(text + "...", True, self.text_color)
surface.blit(text_surface, text_rect)
def handle_event(self, event: pygame.event.Event) -> bool:
"""Handle input events for the tree widget."""
if event.type == pygame.MOUSEBUTTONDOWN:
if event.button == 1: # Left click
# Event position is already converted to local coordinates by HUD
node = self._get_node_at_position_local(event.pos)
if node:
# Check for expand/collapse click
if self._is_click_on_expand_collapse_local(node, event.pos):
node.toggle_expand()
self._update_visible_nodes()
else:
# Selection click
multi_select = pygame.key.get_pressed()[pygame.K_LSHIFT]
self.selection_manager.select_node(node, multi_select)
return True
elif event.type == pygame.MOUSEMOTION:
# Event position is already converted to local coordinates by HUD
self.hovered_node = self._get_node_at_position_local(event.pos)
return True
elif event.type == pygame.MOUSEWHEEL:
# Handle scrolling
if event.y > 0:
self.scroll_offset = max(0, self.scroll_offset - self.node_height)
else:
max_scroll = max(0, self.total_height - self.rect.height)
self.scroll_offset = min(max_scroll, self.scroll_offset + self.node_height)
self._update_node_layout()
return True
return False
def update(self, time_delta: float) -> None:
"""Update the tree widget with smart performance optimizations."""
# Smart update system: only update tree structure when necessary
if self._should_update_tree():
self._update_tree_structure()
self._update_visible_nodes()
def _should_update_tree(self) -> bool:
"""Determine if the tree structure needs updating."""
if not hasattr(self, '_last_entity_count'):
self._last_entity_count = len(self.world.get_objects())
return True
current_count = len(self.world.get_objects())
if current_count != self._last_entity_count:
self._last_entity_count = current_count
return True
# Also check if any previously selected entities have been removed
if self.selection_manager.selected_nodes:
for node in self.selection_manager.selected_nodes[:]:
if hasattr(node, 'entity'):
try:
# Check if entity still exists in world
if node.entity not in self.world.get_objects():
return True
except:
return True
return False
def draw(self, surface: pygame.Surface) -> None:
"""Draw the tree widget."""
# Create a clipping surface for the tree area
tree_surface = pygame.Surface((self.rect.width, self.rect.height))
tree_surface.fill((30, 30, 40)) # Background color
# Set up clipping rect to prevent rendering outside bounds
clip_rect = pygame.Rect(0, 0, self.rect.width, self.rect.height)
tree_surface.set_clip(clip_rect)
# Render visible nodes (only those that should be visible after scrolling)
for node in self.visible_nodes:
# Only render if node is within the visible tree area
if node.rect.bottom > 0 and node.rect.top < self.rect.height:
self._render_node(tree_surface, node)
# Blit to main surface at position (0,0) since this is already positioned
surface.blit(tree_surface, (0, 0))
def expand_all(self) -> None:
"""Expand all nodes in the tree."""
for node in self.all_nodes:
if node.can_expand():
node.expand()
self._update_visible_nodes()
def collapse_all(self) -> None:
"""Collapse all nodes except the root."""
for node in self.all_nodes[1:]: # Skip root node
if node.can_expand():
node.collapse()
self._update_visible_nodes()
def get_selected_entities(self) -> List[Any]:
"""Get the entities currently selected in the tree."""
return self.selection_manager.get_selected_entities()
def select_entities(self, entities: List[Any]) -> None:
"""Select specific entities in the tree."""
self.selection_manager.clear_selection()
for entity in entities:
# Find the entity node for this entity
for node in self.all_nodes:
if hasattr(node, 'entity') and node.entity == entity:
self.selection_manager.select_node(node, multi_select=True)
break
# Expand parent nodes to show selected entities
for node in self.selection_manager.selected_nodes:
parent = node.parent
while parent:
if parent.can_expand():
parent.expand()
parent = parent.parent
self._update_visible_nodes()
def scroll_to_node(self, node: TreeNode) -> None:
"""Scroll the tree to show the given node."""
if node not in self.visible_nodes:
return
node_top = node.rect.y
node_bottom = node.rect.y + node.rect.height
if node_top < 0:
self.scroll_offset = max(0, self.scroll_offset + node_top)
elif node_bottom > self.rect.height:
self.scroll_offset = min(
self.total_height - self.rect.height,
self.scroll_offset + (node_bottom - self.rect.height)
)
self._update_node_layout()