diff --git a/tools/jotpluggler/datatree.py b/tools/jotpluggler/datatree.py index 7bd026cf98..c18ab61892 100644 --- a/tools/jotpluggler/datatree.py +++ b/tools/jotpluggler/datatree.py @@ -2,7 +2,6 @@ import os import re import threading import numpy as np -from collections import deque import dearpygui.dearpygui as dpg @@ -12,8 +11,9 @@ class DataTreeNode: self.full_path = full_path self.parent = parent self.children: dict[str, DataTreeNode] = {} + self.filtered_children: dict[str, DataTreeNode] = {} + self.created_children: dict[str, DataTreeNode] = {} self.is_leaf = False - self.child_count = 0 self.is_plottable: bool | None = None self.ui_created = False self.children_ui_created = False @@ -28,13 +28,17 @@ class DataTree: self.playback_manager = playback_manager self.current_search = "" self.data_tree = DataTreeNode(name="root") - self._build_queue: deque[tuple[DataTreeNode, str | None, str | int]] = deque() - self._all_paths_cache: set[str] = set() - self._item_handlers: set[str] = set() + self._build_queue: dict[str, tuple[DataTreeNode, DataTreeNode, str | int]] = {} # full_path -> (node, parent, before_tag) + self._current_created_paths: set[str] = set() + self._current_filtered_paths: set[str] = set() + self._path_to_node: dict[str, DataTreeNode] = {} # full_path -> node + self._expanded_tags: set[str] = set() + self._item_handlers: dict[str, str] = {} # ui_tag -> handler_tag self._avg_char_width = None self._queued_search = None self._new_data = False self._ui_lock = threading.RLock() + self._handlers_to_delete = [] self.data_manager.add_observer(self._on_data_loaded) def create_ui(self, parent_tag: str): @@ -48,134 +52,185 @@ class DataTree: def _on_data_loaded(self, data: dict): with self._ui_lock: - if data.get('segment_added'): + if data.get('segment_added') or data.get('reset'): self._new_data = True - elif data.get('reset'): - self._all_paths_cache = set() - self._new_data = True - - - def _populate_tree(self): - self._clear_ui() - self.data_tree = self._add_paths_to_tree(self._all_paths_cache, incremental=False) - if self.data_tree: - self._request_children_build(self.data_tree) - - def _add_paths_to_tree(self, paths, incremental=False): - search_term = self.current_search.strip().lower() - filtered_paths = [path for path in paths if self._should_show_path(path, search_term)] - target_tree = self.data_tree if incremental else DataTreeNode(name="root") - - if not filtered_paths: - return target_tree - - parent_nodes_to_recheck = set() - for path in sorted(filtered_paths): - parts = path.split('/') - current_node = target_tree - current_path_prefix = "" - - for i, part in enumerate(parts): - current_path_prefix = f"{current_path_prefix}/{part}" if current_path_prefix else part - if i < len(parts) - 1: - parent_nodes_to_recheck.add(current_node) # for incremental changes from new data - if part not in current_node.children: - current_node.children[part] = DataTreeNode(name=part, full_path=current_path_prefix, parent=current_node) - current_node = current_node.children[part] - - if not current_node.is_leaf: - current_node.is_leaf = True - - self._calculate_child_counts(target_tree) - if incremental: - for p_node in parent_nodes_to_recheck: - p_node.children_ui_created = False - self._request_children_build(p_node) - return target_tree def update_frame(self, font): + if self._handlers_to_delete: # we need to do everything in main thread, frame callbacks are flaky + dpg.render_dearpygui_frame() # wait a frame to ensure queued callbacks are done + with self._ui_lock: + for handler in self._handlers_to_delete: + dpg.delete_item(handler) + self._handlers_to_delete.clear() + with self._ui_lock: if self._avg_char_width is None and dpg.is_dearpygui_running(): self._avg_char_width = self.calculate_avg_char_width(font) if self._new_data: - current_paths = set(self.data_manager.get_all_paths()) - new_paths = current_paths - self._all_paths_cache - all_paths_empty = not self._all_paths_cache - self._all_paths_cache = current_paths - if all_paths_empty: - self._populate_tree() - elif new_paths: - self._add_paths_to_tree(new_paths, incremental=True) + self._process_path_change() self._new_data = False return if self._queued_search is not None: self.current_search = self._queued_search - self._all_paths_cache = set(self.data_manager.get_all_paths()) - self._populate_tree() + self._process_path_change() self._queued_search = None return nodes_processed = 0 while self._build_queue and nodes_processed < self.MAX_NODES_PER_FRAME: - child_node, parent_tag, before_tag = self._build_queue.popleft() + child_node, parent, before_tag = self._build_queue.pop(next(iter(self._build_queue))) + parent_tag = "data_tree_container" if parent.name == "root" else parent.ui_tag if not child_node.ui_created: if child_node.is_leaf: self._create_leaf_ui(child_node, parent_tag, before_tag) else: self._create_tree_node_ui(child_node, parent_tag, before_tag) + parent.created_children[child_node.name] = parent.children[child_node.name] + self._current_created_paths.add(child_node.full_path) nodes_processed += 1 - def search_data(self): - self._queued_search = dpg.get_value("search_input") + def _process_path_change(self): + self._build_queue.clear() + search_term = self.current_search.strip().lower() + all_paths = set(self.data_manager.get_all_paths()) + new_filtered_leafs = {path for path in all_paths if self._should_show_path(path, search_term)} + new_filtered_paths = set(new_filtered_leafs) + for path in new_filtered_leafs: + parts = path.split('/') + for i in range(1, len(parts)): + prefix = '/'.join(parts[:i]) + new_filtered_paths.add(prefix) + created_paths_to_remove = self._current_created_paths - new_filtered_paths + filtered_paths_to_remove = self._current_filtered_paths - new_filtered_leafs + + if created_paths_to_remove or filtered_paths_to_remove: + self._remove_paths_from_tree(created_paths_to_remove, filtered_paths_to_remove) + self._apply_expansion_to_tree(self.data_tree, search_term) + + paths_to_add = new_filtered_leafs - self._current_created_paths + if paths_to_add: + self._add_paths_to_tree(paths_to_add) + self._apply_expansion_to_tree(self.data_tree, search_term) + self._current_filtered_paths = new_filtered_paths + + def _remove_paths_from_tree(self, created_paths_to_remove, filtered_paths_to_remove): + for path in sorted(created_paths_to_remove, reverse=True): + current_node = self._path_to_node[path] + + if len(current_node.created_children) == 0: + self._current_created_paths.remove(current_node.full_path) + if item_handler_tag := self._item_handlers.get(current_node.ui_tag): + dpg.configure_item(item_handler_tag, show=False) + self._handlers_to_delete.append(item_handler_tag) + del self._item_handlers[current_node.ui_tag] + dpg.delete_item(current_node.ui_tag) + current_node.ui_created = False + current_node.ui_tag = None + current_node.children_ui_created = False + del current_node.parent.created_children[current_node.name] + del current_node.parent.filtered_children[current_node.name] + + for path in filtered_paths_to_remove: + parts = path.split('/') + current_node = self._path_to_node[path] - def _clear_ui(self): - for handler_tag in self._item_handlers: - dpg.configure_item(handler_tag, show=False) - dpg.set_frame_callback(dpg.get_frame_count() + 1, callback=self._delete_handlers, user_data=list(self._item_handlers)) - self._item_handlers.clear() + part_array_index = -1 + while len(current_node.filtered_children) == 0 and part_array_index >= -len(parts): + current_node = current_node.parent + if parts[part_array_index] in current_node.filtered_children: + del current_node.filtered_children[parts[part_array_index]] + part_array_index -= 1 - if dpg.does_item_exist("data_tree_container"): - dpg.delete_item("data_tree_container", children_only=True) + def _add_paths_to_tree(self, paths): + parent_nodes_to_recheck = set() + for path in sorted(paths): + parts = path.split('/') + current_node = self.data_tree + current_path_prefix = "" - self._build_queue.clear() + for i, part in enumerate(parts): + current_path_prefix = f"{current_path_prefix}/{part}" if current_path_prefix else part + if i < len(parts) - 1: + parent_nodes_to_recheck.add(current_node) # for incremental changes from new data + if part not in current_node.children: + current_node.children[part] = DataTreeNode(name=part, full_path=current_path_prefix, parent=current_node) + self._path_to_node[current_path_prefix] = current_node.children[part] + current_node.filtered_children[part] = current_node.children[part] + current_node = current_node.children[part] + + if not current_node.is_leaf: + current_node.is_leaf = True + + for p_node in parent_nodes_to_recheck: + p_node.children_ui_created = False + self._request_children_build(p_node) - def _delete_handlers(self, sender, app_data, user_data): - for handler in user_data: - dpg.delete_item(handler) + def _get_node_label_and_expand(self, node: DataTreeNode, search_term: str): + label = f"{node.name} ({len(node.filtered_children)} fields)" + expand = len(search_term) > 0 and any(search_term in path for path in self._get_descendant_paths(node)) + if expand and node.parent and len(node.parent.filtered_children) > 100 and len(node.filtered_children) > 2: + label += " (+)" # symbol for large lists which aren't fully expanded for performance (only affects procLog rn) + expand = False + return label, expand + + def _apply_expansion_to_tree(self, node: DataTreeNode, search_term: str): + if node.ui_created and not node.is_leaf and node.ui_tag and dpg.does_item_exist(node.ui_tag): + label, expand = self._get_node_label_and_expand(node, search_term) + if expand: + self._expanded_tags.add(node.ui_tag) + dpg.set_value(node.ui_tag, expand) + elif node.ui_tag in self._expanded_tags: # not expanded and was expanded + self._expanded_tags.remove(node.ui_tag) + dpg.set_value(node.ui_tag, expand) + dpg.delete_item(node.ui_tag, children_only=True) # delete children (not visible since collapsed) + self._reset_ui_state_recursive(node) + node.children_ui_created = False + dpg.set_item_label(node.ui_tag, label) + for child in node.created_children.values(): + self._apply_expansion_to_tree(child, search_term) + + def _reset_ui_state_recursive(self, node: DataTreeNode): + for child in node.created_children.values(): + if child.ui_tag is not None: + if item_handler_tag := self._item_handlers.get(child.ui_tag): + self._handlers_to_delete.append(item_handler_tag) + dpg.configure_item(item_handler_tag, show=False) + del self._item_handlers[child.ui_tag] + self._reset_ui_state_recursive(child) + child.ui_created = False + child.ui_tag = None + child.children_ui_created = False + self._current_created_paths.remove(child.full_path) + node.created_children.clear() - def _calculate_child_counts(self, node: DataTreeNode): - if node.is_leaf: - node.child_count = 0 - else: - node.child_count = len(node.children) - for child in node.children.values(): - self._calculate_child_counts(child) + def search_data(self): + with self._ui_lock: + self._queued_search = dpg.get_value("search_input") def _create_tree_node_ui(self, node: DataTreeNode, parent_tag: str, before: str | int): - tag = f"tree_{node.full_path}" - node.ui_tag = tag - label = f"{node.name} ({node.child_count} fields)" + node.ui_tag = f"tree_{node.full_path}" search_term = self.current_search.strip().lower() - expand = bool(search_term) and len(search_term) > 1 and any(search_term in path for path in self._get_descendant_paths(node)) - if expand and node.parent and node.parent.child_count > 100 and node.child_count > 2: # don't fully autoexpand large lists (only affects procLog rn) - label += " (+)" - expand = False + label, expand = self._get_node_label_and_expand(node, search_term) + if expand: + self._expanded_tags.add(node.ui_tag) + elif node.ui_tag in self._expanded_tags: + self._expanded_tags.remove(node.ui_tag) with dpg.tree_node( - label=label, parent=parent_tag, tag=tag, default_open=expand, open_on_arrow=True, open_on_double_click=True, before=before, delay_search=True + label=label, parent=parent_tag, tag=node.ui_tag, default_open=expand, open_on_arrow=True, open_on_double_click=True, before=before, delay_search=True ): with dpg.item_handler_registry() as handler_tag: dpg.add_item_toggled_open_handler(callback=lambda s, a, u: self._request_children_build(node)) dpg.add_item_visible_handler(callback=lambda s, a, u: self._request_children_build(node)) - dpg.bind_item_handler_registry(tag, handler_tag) - self._item_handlers.add(handler_tag) - + dpg.bind_item_handler_registry(node.ui_tag, handler_tag) + self._item_handlers[node.ui_tag] = handler_tag node.ui_created = True def _create_leaf_ui(self, node: DataTreeNode, parent_tag: str, before: str | int): - with dpg.group(parent=parent_tag, tag=f"leaf_{node.full_path}", before=before, delay_search=True) as draggable_group: + node.ui_tag = f"leaf_{node.full_path}" + with dpg.group(parent=parent_tag, tag=node.ui_tag, before=before, delay_search=True): with dpg.table(header_row=False, policy=dpg.mvTable_SizingStretchProp, delay_search=True): dpg.add_table_column(init_width_or_weight=0.5) dpg.add_table_column(init_width_or_weight=0.5) @@ -186,21 +241,21 @@ class DataTree: if node.is_plottable is None: node.is_plottable = self.data_manager.is_plottable(node.full_path) if node.is_plottable: - with dpg.drag_payload(parent=draggable_group, drag_data=node.full_path, payload_type="TIMESERIES_PAYLOAD"): + with dpg.drag_payload(parent=node.ui_tag, drag_data=node.full_path, payload_type="TIMESERIES_PAYLOAD"): dpg.add_text(f"Plot: {node.full_path}") with dpg.item_handler_registry() as handler_tag: dpg.add_item_visible_handler(callback=self._on_item_visible, user_data=node.full_path) - dpg.bind_item_handler_registry(draggable_group, handler_tag) - self._item_handlers.add(handler_tag) - + dpg.bind_item_handler_registry(node.ui_tag, handler_tag) + self._item_handlers[node.ui_tag] = handler_tag node.ui_created = True - node.ui_tag = f"value_{node.full_path}" def _on_item_visible(self, sender, app_data, user_data): with self._ui_lock: path = user_data value_tag = f"value_{path}" + if not dpg.does_item_exist(value_tag): + return value_column_width = dpg.get_item_rect_size("sidebar_window")[0] // 2 value = self.data_manager.get_value_at(path, self.playback_manager.current_time_s) if value is not None: @@ -212,8 +267,7 @@ class DataTree: def _request_children_build(self, node: DataTreeNode): with self._ui_lock: if not node.children_ui_created and (node.name == "root" or (node.ui_tag is not None and dpg.get_value(node.ui_tag))): # check root or node expanded - parent_tag = "data_tree_container" if node.name == "root" else node.ui_tag - sorted_children = sorted(node.children.values(), key=self._natural_sort_key) + sorted_children = sorted(node.filtered_children.values(), key=self._natural_sort_key) next_existing: list[int | str] = [0] * len(sorted_children) current_before_tag: int | str = 0 @@ -228,7 +282,7 @@ class DataTree: for i, child_node in enumerate(sorted_children): if not child_node.ui_created: before_tag = next_existing[i] - self._build_queue.append((child_node, parent_tag, before_tag)) + self._build_queue[child_node.full_path] = (child_node, node, before_tag) node.children_ui_created = True def _should_show_path(self, path: str, search_term: str) -> bool: @@ -242,7 +296,7 @@ class DataTree: return (node_type_key, parts) def _get_descendant_paths(self, node: DataTreeNode): - for child_name, child_node in node.children.items(): + for child_name, child_node in node.filtered_children.items(): child_name_lower = child_name.lower() if child_node.is_leaf: yield child_name_lower