yes

Signed-off-by: Max Chesterfield <max.chesterfield@zepben.com>

Author: chestm007

changelog.md

@@ -92,7 +92,8 @@
 * None.
 
 ### Enhancements
-* None.
+* `EquipmentTreeBuilder` now has a `leaves` attribute which is a dict of all mrid -> IdentifiedObject contained in the tree
+* `TreeNode` uses a set to store `_children` instead of a set, massive performance gains.
 
 ### Fixes
 * Marked some extensions properties and classes with [ZBEX] that were missing them (might still be more). In addition to the ones moved into the extensions

src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py

@@ -35,6 +35,7 @@ class EquipmentTreeBuilder(StepActionWithContextValue):
     """
 
     _roots: dict[ConductingEquipment, EquipmentTreeNode] = {}
+    leaves: dict[str, EquipmentTreeNode] = {}
 
     def __init__(self):
         super().__init__(key=str(uuid.uuid4()))
@@ -66,6 +67,7 @@ def _apply(self, item: NetworkTraceStep[Any], context: StepContext):
         current_node: TreeNode = self.get_context_value(context)
         if current_node.parent:
             current_node.parent.add_child(current_node)
+        self.leaves[current_node.identified_object.mrid] = current_node
 
     def clear(self):
         self._roots.clear()

src/zepben/ewb/services/network/tracing/networktrace/actions/tree_node.py

@@ -5,7 +5,7 @@
 
 __all__ = ['TreeNode']
 
-from typing import List, TypeVar, Generic
+from typing import List, TypeVar, Generic, Set
 
 from zepben.ewb import IdentifiedObject
 
@@ -20,7 +20,7 @@ class TreeNode(Generic[T]):
     def __init__(self, identified_object: IdentifiedObject, parent=None):
         self.identified_object = identified_object
         self._parent: TreeNode = parent
-        self._children: List[TreeNode] = []
+        self._children: Set[TreeNode] = set()
 
     @property
     def parent(self) -> 'TreeNode[T]':
@@ -31,7 +31,7 @@ def children(self):
         return list(self._children)
 
     def add_child(self, child: 'TreeNode'):
-        self._children.append(child)
+        self._children.add(child)
 
     def __str__(self):
         return f"{{object: {self.identified_object}, parent: {self.parent or ''}, num children: {len(self.children)}}}"

test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py

@@ -54,9 +54,9 @@ async def test_downstream_tree():
     root = list(tree_builder.roots)[0]
 
     assert root is not None
-    _verify_tree_asset(root, n["j1"], None, [n["ac1"], n["ac3"]])
+    _verify_tree_asset(root, n["j1"], None, [n["ac3"], n["ac1"]])
 
-    test_node = root.children[0]
+    test_node = root.children[1]
     _verify_tree_asset(test_node, n["ac1"], n["j1"], [n["j2"]])
 
     test_node = test_node.children[0]
@@ -74,7 +74,7 @@ async def test_downstream_tree():
     test_node = next(iter(test_node.children))
     _verify_tree_asset(test_node, n["j6"], n["ac4"], [])
 
-    test_node = list(root.children)[1]
+    test_node = list(root.children)[0]
     _verify_tree_asset(test_node, n["ac3"], n["j1"], [n["j4"]])
 
     test_node = next(iter(test_node.children))