320 lines
12 KiB
Python
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() |