From 35622e4394319826250a698f4695b2bb7b307ef9 Mon Sep 17 00:00:00 2001 From: Max Chesterfield Date: Tue, 29 Jul 2025 14:44:17 +1000 Subject: [PATCH 1/4] leaves added, tests fixed Signed-off-by: Max Chesterfield --- .../actions/equipment_tree_builder.py | 23 +++++++ .../tracing/networktrace/actions/tree_node.py | 18 +++-- .../actions/test_equipment_tree_builder.py | 67 ++++++++++++------- 3 files changed, 74 insertions(+), 34 deletions(-) diff --git a/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py b/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py index 18d7f5989..a441d3135 100644 --- a/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py +++ b/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py @@ -35,6 +35,7 @@ class EquipmentTreeBuilder(StepActionWithContextValue): """ _roots: dict[ConductingEquipment, EquipmentTreeNode] = {} + _leaves: set[EquipmentTreeNode] = set() def __init__(self): super().__init__(key=str(uuid.uuid4())) @@ -43,6 +44,26 @@ def __init__(self): def roots(self) -> Generator[TreeNode[ConductingEquipment], None, None]: return (r for r in self._roots.values()) + def recurse_nodes(self) -> Generator[TreeNode[ConductingEquipment], None, None]: + """ + Returns a generator that will yield every node in the tree structure. + """ + def recurse(node: TreeNode[ConductingEquipment]): + yield node + for child in node.children: + yield from recurse(child) + + for root in self._roots.values(): + yield from recurse(root) + + @property + def leaves(self) -> set[EquipmentTreeNode]: + """ + Return the leaves of the tree structure. Depending on how the backing trace is configured, + there may be extra unexpected leaves in loops. + """ + return set(self._leaves) + def compute_initial_value(self, item: NetworkTraceStep[Any]) -> EquipmentTreeNode: node = self._roots.get(item.path.to_equipment) if node is None: @@ -64,7 +85,9 @@ def compute_next_value( def _apply(self, item: NetworkTraceStep[Any], context: StepContext): current_node: TreeNode = self.get_context_value(context) + self._leaves.add(current_node) # add this node to _leaves as it has no children if current_node.parent: + self._leaves.discard(current_node.parent) # this nodes parent now has a child, it's not a leaf anymore current_node.parent.add_child(current_node) def clear(self): diff --git a/src/zepben/ewb/services/network/tracing/networktrace/actions/tree_node.py b/src/zepben/ewb/services/network/tracing/networktrace/actions/tree_node.py index eb8ddae69..f4587eb61 100644 --- a/src/zepben/ewb/services/network/tracing/networktrace/actions/tree_node.py +++ b/src/zepben/ewb/services/network/tracing/networktrace/actions/tree_node.py @@ -5,33 +5,31 @@ __all__ = ['TreeNode'] -from typing import List, TypeVar, Generic - -from zepben.ewb import IdentifiedObject +from typing import TypeVar, Generic, Set T = TypeVar('T') class TreeNode(Generic[T]): """ - represents a node in the NetworkTrace tree + Represents a node in the NetworkTrace tree """ - def __init__(self, identified_object: IdentifiedObject, parent=None): + def __init__(self, identified_object: T, 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]': return self._parent @property - def children(self): - return list(self._children) + def children(self) -> Set['TreeNode[T]']: + return set(self._children) - def add_child(self, child: 'TreeNode'): - self._children.append(child) + def add_child(self, child: 'TreeNode[T]'): + self._children.add(child) def __str__(self): return f"{{object: {self.identified_object}, parent: {self.parent or ''}, num children: {len(self.children)}}}" diff --git a/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py b/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py index e9341d177..e10b0e978 100644 --- a/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py +++ b/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py @@ -33,12 +33,15 @@ async def test_downstream_tree(): start = n.get("j1", ConductingEquipment) assert start is not None tree_builder = EquipmentTreeBuilder() - trace = Tracing.network_trace_branching( - network_state_operators=normal, - action_step_type=NetworkTraceActionType.FIRST_STEP_ON_EQUIPMENT) \ - .add_condition(downstream()) \ - .add_step_action(tree_builder) \ + trace = ( + Tracing.network_trace_branching( + network_state_operators=normal, + action_step_type=NetworkTraceActionType.FIRST_STEP_ON_EQUIPMENT + ) + .add_condition(downstream()) + .add_step_action(tree_builder) .add_step_action(lambda item, context: visited_ce.append(item.path.to_equipment.mrid)) + ) await trace.run(start) @@ -51,34 +54,39 @@ async def test_downstream_tree(): pprint.pprint(visit_counts) - root = list(tree_builder.roots)[0] + root = tree_builder._roots[start] assert root is not None _verify_tree_asset(root, n["j1"], None, [n["ac1"], n["ac3"]]) - test_node = root.children[0] - _verify_tree_asset(test_node, n["ac1"], n["j1"], [n["j2"]]) + assert len(root.children) == 2 + for test_node in root.children: + if test_node.identified_object == n['ac1']: + _verify_tree_asset(test_node, n["ac1"], n["j1"], [n["j2"]]) - test_node = test_node.children[0] - _verify_tree_asset(test_node, n["j2"], n["ac1"], [n["ac2"]]) + test_node = test_node.children.pop() + _verify_tree_asset(test_node, n["j2"], n["ac1"], [n["ac2"]]) - test_node = test_node.children[0] - _verify_tree_asset(test_node, n["ac2"], n["j2"], [n["j3"]]) + test_node = test_node.children.pop() + _verify_tree_asset(test_node, n["ac2"], n["j2"], [n["j3"]]) - test_node = next(iter(test_node.children)) - _verify_tree_asset(test_node, n["j3"], n["ac2"], [n["ac4"]]) + test_node = next(iter(test_node.children)) + _verify_tree_asset(test_node, n["j3"], n["ac2"], [n["ac4"]]) - test_node = next(iter(test_node.children)) - _verify_tree_asset(test_node, n["ac4"], n["j3"], [n["j6"]]) + test_node = next(iter(test_node.children)) + _verify_tree_asset(test_node, n["ac4"], n["j3"], [n["j6"]]) - test_node = next(iter(test_node.children)) - _verify_tree_asset(test_node, n["j6"], n["ac4"], []) + test_node = next(iter(test_node.children)) + _verify_tree_asset(test_node, n["j6"], n["ac4"], []) + break - test_node = list(root.children)[1] - _verify_tree_asset(test_node, n["ac3"], n["j1"], [n["j4"]]) + elif test_node.identified_object == n['ac3']: + _verify_tree_asset(test_node, n["ac3"], n["j1"], [n["j4"]]) - test_node = next(iter(test_node.children)) - _verify_tree_asset(test_node, n["j4"], n["ac3"], [n["ac5"], n["ac6"]]) + test_node = next(iter(test_node.children)) + _verify_tree_asset(test_node, n["j4"], n["ac3"], [n["ac5"], n["ac6"]]) + else: + assert False assert len(_find_nodes(root, "j0")) == 0 assert len(_find_nodes(root, "ac0")) == 0 @@ -147,6 +155,10 @@ async def test_downstream_tree(): assert _find_node_depths(root, "ac16") == [8, 9, 11, 14] + for ce in (n['j5'], n['j13']): + assert ce in {l.identified_object for l in tree_builder.leaves} + + def _verify_tree_asset( tree_node: TreeNode, expected_asset: Optional[ConductingEquipment], @@ -162,8 +174,15 @@ def _verify_tree_asset( else: assert tree_node.parent is None - children_nodes = list(c.identified_object for c in tree_node.children) - assert children_nodes == expected_children + children_nodes = [c.identified_object for c in tree_node.children] + try: + for child in expected_children: + assert child in children_nodes + for child in children_nodes: + assert child in expected_children + except AssertionError as e: + e.args = (expected_children, children_nodes) + raise e def _find_nodes(root: TreeNode[ConductingEquipment], asset_id: str) -> List[TreeNode[ConductingEquipment]]: From e74c68f96971f1b77efebcaba0914495ff18d7ae Mon Sep 17 00:00:00 2001 From: Max Chesterfield Date: Mon, 4 Aug 2025 15:23:39 +1000 Subject: [PATCH 2/4] leaves now calculated when specified Signed-off-by: Max Chesterfield --- .../actions/equipment_tree_builder.py | 26 +++++++--- .../actions/test_equipment_tree_builder.py | 39 +++++++++++++-- .../networktrace/test_network_trace.py | 50 ++++++++++++++++++- 3 files changed, 101 insertions(+), 14 deletions(-) diff --git a/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py b/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py index a441d3135..217e7ff09 100644 --- a/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py +++ b/src/zepben/ewb/services/network/tracing/networktrace/actions/equipment_tree_builder.py @@ -34,12 +34,14 @@ class EquipmentTreeBuilder(StepActionWithContextValue): >>> .add_step_action(tree_builder)).run() """ - _roots: dict[ConductingEquipment, EquipmentTreeNode] = {} - _leaves: set[EquipmentTreeNode] = set() - - def __init__(self): + def __init__(self, calculate_leaves: bool = False): super().__init__(key=str(uuid.uuid4())) + self._roots: dict[ConductingEquipment, EquipmentTreeNode] = {} + self._leaves: set[EquipmentTreeNode] = set() + + self._calculate_leaves = calculate_leaves + @property def roots(self) -> Generator[TreeNode[ConductingEquipment], None, None]: return (r for r in self._roots.values()) @@ -59,9 +61,11 @@ def recurse(node: TreeNode[ConductingEquipment]): @property def leaves(self) -> set[EquipmentTreeNode]: """ - Return the leaves of the tree structure. Depending on how the backing trace is configured, - there may be extra unexpected leaves in loops. + Return the leaves of the tree structure. + Depending on how the backing trace is configured, there may be extra unexpected leaves in loops. """ + if not self._calculate_leaves: + raise AttributeError('leaves were not calculated, you must pass calculate_leaves=True to the EquipmentTreeBuilder when creating.') return set(self._leaves) def compute_initial_value(self, item: NetworkTraceStep[Any]) -> EquipmentTreeNode: @@ -83,12 +87,18 @@ def compute_next_value( else: return TreeNode(next_item.path.to_equipment, current_value) - def _apply(self, item: NetworkTraceStep[Any], context: StepContext): + def _apply(self, _: NetworkTraceStep[Any], context: StepContext): current_node: TreeNode = self.get_context_value(context) + if current_node.parent: + current_node.parent.add_child(current_node) + + if self._calculate_leaves: + self._process_leaf(current_node) + + def _process_leaf(self, current_node: TreeNode[ConductingEquipment]): self._leaves.add(current_node) # add this node to _leaves as it has no children if current_node.parent: self._leaves.discard(current_node.parent) # this nodes parent now has a child, it's not a leaf anymore - current_node.parent.add_child(current_node) def clear(self): self._roots.clear() diff --git a/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py b/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py index e10b0e978..f104a42c5 100644 --- a/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py +++ b/test/services/network/tracing/networktrace/actions/test_equipment_tree_builder.py @@ -16,6 +16,41 @@ from zepben.ewb.services.network.tracing.networktrace.actions.tree_node import TreeNode +def test_accessing_leaves_when_not_calculated_raises_exception(): + builder = EquipmentTreeBuilder() + with pytest.raises(AttributeError): + builder.leaves + +@pytest.mark.asyncio +async def test_equipment_tree_builder_leaves(): + n = create_looping_network() + normal = NetworkStateOperators.NORMAL + current = NetworkStateOperators.CURRENT + + await Tracing.set_phases().run(n) + feeder_head = n.get("j0", ConductingEquipment) + await Tracing.set_direction().run_terminal(feeder_head, network_state_operators=normal) + await Tracing.set_direction().run_terminal(feeder_head, network_state_operators=current) + await log_directions(n.get('j0', ConductingEquipment)) + + start = n.get("j1", ConductingEquipment) + assert start is not None + tree_builder = EquipmentTreeBuilder(calculate_leaves=True) + trace = ( + Tracing.network_trace_branching( + network_state_operators=normal, + action_step_type=NetworkTraceActionType.FIRST_STEP_ON_EQUIPMENT + ) + .add_condition(downstream()) + .add_step_action(tree_builder) + ) + + await trace.run(start) + + for ce in (n['j5'], n['j13']): + assert ce in {l.identified_object for l in tree_builder.leaves} + + @pytest.mark.asyncio async def test_downstream_tree(): n = create_looping_network() @@ -155,10 +190,6 @@ async def test_downstream_tree(): assert _find_node_depths(root, "ac16") == [8, 9, 11, 14] - for ce in (n['j5'], n['j13']): - assert ce in {l.identified_object for l in tree_builder.leaves} - - def _verify_tree_asset( tree_node: TreeNode, expected_asset: Optional[ConductingEquipment], diff --git a/test/services/network/tracing/networktrace/test_network_trace.py b/test/services/network/tracing/networktrace/test_network_trace.py index 7d162e3e3..ba97d54a8 100644 --- a/test/services/network/tracing/networktrace/test_network_trace.py +++ b/test/services/network/tracing/networktrace/test_network_trace.py @@ -13,8 +13,8 @@ from services.network.tracing.networktrace.test_network_trace_step_path_provider import PathTerminal, _verify_paths from zepben.ewb import AcLineSegment, Clamp, Terminal, NetworkTraceStep, Cut, ConductingEquipment, TraversalQueue, Junction, ngen, NetworkTraceActionType, \ - Tracing -from zepben.ewb.testing.test_network_builder import TestNetworkBuilder + Tracing, StepContext, StepActionWithContextValue, EnergyConsumer +from zepben.evolve.testing.test_network_builder import TestNetworkBuilder Terminal.__add__ = PathTerminal.__add__ Terminal.__sub__ = PathTerminal.__sub__ @@ -286,3 +286,49 @@ async def validate(start: Tuple[str, str], action_step_type: NetworkTraceActionT # Can even use bizarre paths, they are just the same as any other external path. await validate(('c0-t1', 'c2-t1'), NetworkTraceActionType.ALL_STEPS, ["c2-t1", "c2-t2"]) await validate(('c0-t1', 'c2-t1'), NetworkTraceActionType.FIRST_STEP_ON_EQUIPMENT, ["c2-t1"]) + + async def test_context_is_not_shared_between_branches(self): + # + # 1--c0--21--c1-*-21--c2--21--ec3 + # 1 + # 1--c4--21--ec5 + # + + ns = ( + TestNetworkBuilder() + .from_acls() # c0 + .to_acls() # c1 + .with_clamp() # c1-clamp1 + .to_acls() # c2 + .to_energy_consumer() # ec3 + .branch_from('c1-clamp1') + .to_acls() # c4 + .to_energy_consumer() # ec5 + ).network + + data_capture: list[tuple[str, list]] = [] + + class StepActionWithContext(StepActionWithContextValue): + def _apply(self, item: NetworkTraceStep, context: List): + print(item) + if isinstance((ec := item.path.to_equipment), EnergyConsumer): + data_capture.append((ec.mrid, self.get_context_value(context))) + + def compute_next_value(self, next_item: NetworkTraceStep, current_item: NetworkTraceStep, current_value: List): + nv = list(current_value) + nv.append(next_item.path.from_equipment.mrid) + return nv + + def compute_initial_value(self, item: NetworkTraceStep): + return [item.path.from_equipment.mrid] + + await ( + Tracing.network_trace() + .add_step_action(StepActionWithContext('key')) + ).run(ns.get('c0')) + + assert len(data_capture) == 2 + assert data_capture == [('ec3', ['c0', 'c0', 'c1', 'c1', 'c2', 'c2']), + ('ec5', ['c0', 'c0', 'c1', 'c1-clamp1', 'c4', 'c4'])] + + From 9e95937804930f2ca2bf8b4db9e0401cf75285de Mon Sep 17 00:00:00 2001 From: Max Chesterfield Date: Mon, 4 Aug 2025 15:24:15 +1000 Subject: [PATCH 3/4] changelog Signed-off-by: Max Chesterfield --- changelog.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/changelog.md b/changelog.md index e9f9f998b..3f6ae7a67 100644 --- a/changelog.md +++ b/changelog.md @@ -92,7 +92,7 @@ * None. ### Enhancements -* None. +* `EquipmentTreeBuilder` will now calculate `leaves` when specified to do so. ### 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 @@ -127,6 +127,8 @@ * `RegulatingControl.ratedCurrent` * `Sensor.relayFunctions` * `UsagePoint.approvedInverterCapacity` +* using `EquipmentTreeBuilder` more then once per interpreter will no longer cause the `roots` to contain more objects then it should due to `_roots` being a + class var ### Notes * None. From 64a2ec957d3857682278ed978cce589b545cc6a3 Mon Sep 17 00:00:00 2001 From: Max Chesterfield Date: Mon, 4 Aug 2025 15:26:42 +1000 Subject: [PATCH 4/4] huh? Signed-off-by: Max Chesterfield --- .../network/tracing/networktrace/test_network_trace.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/services/network/tracing/networktrace/test_network_trace.py b/test/services/network/tracing/networktrace/test_network_trace.py index ba97d54a8..0f04d259c 100644 --- a/test/services/network/tracing/networktrace/test_network_trace.py +++ b/test/services/network/tracing/networktrace/test_network_trace.py @@ -13,8 +13,8 @@ from services.network.tracing.networktrace.test_network_trace_step_path_provider import PathTerminal, _verify_paths from zepben.ewb import AcLineSegment, Clamp, Terminal, NetworkTraceStep, Cut, ConductingEquipment, TraversalQueue, Junction, ngen, NetworkTraceActionType, \ - Tracing, StepContext, StepActionWithContextValue, EnergyConsumer -from zepben.evolve.testing.test_network_builder import TestNetworkBuilder + Tracing, StepActionWithContextValue, EnergyConsumer +from zepben.ewb.testing.test_network_builder import TestNetworkBuilder Terminal.__add__ = PathTerminal.__add__ Terminal.__sub__ = PathTerminal.__sub__