zepben.ewb 1.0.0__py3-none-any.whl

zepben/ewb/services/network/tracing/find_swer_equipment.py

@@ -0,0 +1,175 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+from __future__ import annotations
+
+__all__ = ["FindSwerEquipment"]
+
+from typing import Set, Union, AsyncGenerator, Type, TYPE_CHECKING
+
+from typing_extensions import TypeVar
+
+from zepben.ewb.model.cim.iec61970.base.core.conducting_equipment import ConductingEquipment
+from zepben.ewb.model.cim.iec61970.base.core.feeder import Feeder
+from zepben.ewb.model.cim.iec61970.base.core.terminal import Terminal
+from zepben.ewb.model.cim.iec61970.base.wires.power_transformer import PowerTransformer
+from zepben.ewb.model.cim.iec61970.base.wires.switch import Switch
+from zepben.ewb.services.network.network_service import NetworkService
+from zepben.ewb.services.network.tracing.networktrace.conditions.conditions import stop_at_open
+from zepben.ewb.services.network.tracing.networktrace.network_trace import NetworkTrace
+from zepben.ewb.services.network.tracing.networktrace.operators.network_state_operators import NetworkStateOperators
+from zepben.ewb.services.network.tracing.networktrace.tracing import Tracing
+
+if TYPE_CHECKING:
+    from logging import Logger
+
+T = TypeVar
+
+
+class FindSwerEquipment:
+    """
+    A class which can be used for finding the SWER equipment in a [NetworkService] or [Feeder].
+    """
+
+    def __init__(self, debug_logger: Logger = None):
+        self._debug_logger = debug_logger
+
+    async def find(
+        self,
+        to_process: Union[NetworkService, Feeder],
+        network_state_operators: Type[NetworkStateOperators] = NetworkStateOperators.NORMAL
+    ) -> Set[ConductingEquipment]:
+        """
+        Convenience method to call out to `find_all` or `find_on_feeder` based on the class type of `to_process`
+
+        :param to_process: the object to process
+        :param network_state_operators: The `NetworkStateOperators` to be used when finding SWER equipment
+
+        :return: A `Set` of `ConductingEquipment` on `Feeder` that is SWER, or energised via SWER.
+        """
+
+        if isinstance(to_process, Feeder):
+            return set(await self.find_on_feeder(to_process, network_state_operators))
+        elif isinstance(to_process, NetworkService):
+            return set([item async for item in self.find_all(to_process, network_state_operators)])
+        else:
+            raise NotImplementedError
+
+    async def find_all(
+        self,
+        network_service: NetworkService,
+        network_state_operators: Type[NetworkStateOperators] = NetworkStateOperators.NORMAL
+    ) -> AsyncGenerator[ConductingEquipment, None]:
+        """
+        Find the `ConductingEquipment` on any `Feeder` in a `NetworkService` which is SWER. This will include any equipment on the LV network that is energised
+        via SWER.
+
+        :param network_service: The `NetworkService` to process.
+        :param network_state_operators: The `NetworkStateOperators` to be used when finding SWER equipment
+
+        :return: A `Set` of `ConductingEquipment` on `Feeder` that is SWER, or energised via SWER.
+        """
+
+        for feeder in network_service.objects(Feeder):
+            for item in await self.find_on_feeder(feeder, network_state_operators):
+                yield item
+
+    async def find_on_feeder(
+        self,
+        feeder: Feeder,
+        network_state_operators: Type[NetworkStateOperators] = NetworkStateOperators.NORMAL
+    ) -> Set[ConductingEquipment]:
+        """
+        Find the `ConductingEquipment` on a `Feeder` which is SWER. This will include any equipment on the LV network that is energised via SWER.
+
+        :param feeder: The `Feeder` to process.
+        :param network_state_operators: The `NetworkStateOperators` to be used when finding SWER equipment
+
+        :return: A `Set` of `ConductingEquipment` on `feeder` that is SWER, or energised via SWER.
+        """
+
+        swer_equipment: Set[ConductingEquipment] = set()
+
+        # We will add all the SWER transformers to the swer_equipment list before starting any traces to prevent tracing though them by accident. In
+        #  order to do this, we collect the sequence to a list to change the iteration order.
+        for equipment in network_state_operators.get_equipment(feeder):
+            if isinstance(equipment, PowerTransformer):
+                if _has_swer_terminal(equipment) and _has_non_swer_terminal(equipment):
+                    swer_equipment.add(equipment)
+                    await self._trace_from(network_state_operators, equipment, swer_equipment)
+        return swer_equipment
+
+    def _create_trace(self, state_operators: Type[NetworkStateOperators]) -> NetworkTrace[T]:
+        return Tracing.network_trace(
+            network_state_operators=state_operators,
+            debug_logger=self._debug_logger,
+            name=f'FindSwerEquipment({state_operators.description})'
+        ).add_condition(stop_at_open())
+
+    async def _trace_from(self, state_operators: Type[NetworkStateOperators], transformer: PowerTransformer, swer_equipment: Set[ConductingEquipment]):
+        # Trace from any SWER terminals.
+        await self._trace_swer_from(state_operators, transformer, swer_equipment)
+
+        # Trace from any LV terminals.
+        await self._trace_lv_from(state_operators, transformer, swer_equipment)
+
+    async def _trace_swer_from(
+        self,
+        state_operators: Type[NetworkStateOperators],
+        transformer: PowerTransformer,
+        swer_equipment: Set[ConductingEquipment]
+    ):
+
+        def condition(next_step, nctx, step, ctx):
+            if _is_swer_terminal(next_step.path.to_terminal) or isinstance(next_step.path.to_equipment, Switch):
+                return next_step.path.to_equipment not in swer_equipment
+
+        trace = (
+            self._create_trace(state_operators)
+            .add_queue_condition(condition)
+            .add_step_action(lambda step, ctx: swer_equipment.add(step.path.to_equipment))
+        )
+
+        for it in (t for t in transformer.terminals if _is_swer_terminal(t)):
+            trace.reset()
+            await trace.run(it, None)
+
+    async def _trace_lv_from(
+        self,
+        state_operators: Type[NetworkStateOperators],
+        transformer: PowerTransformer,
+        swer_equipment: Set[ConductingEquipment]
+    ):
+
+        def condition(next_step, nctx, step, ctx):
+            if 1 <= next_step.path.to_equipment.base_voltage_value <= 1000:
+                return next_step.path.to_equipment not in swer_equipment
+
+        trace = (
+            self._create_trace(state_operators)
+            .add_queue_condition(condition)
+            .add_step_action(lambda step, ctx: swer_equipment.add(step.path.to_equipment))
+        )
+
+        for terminal in transformer.terminals:
+            if _is_non_swer_terminal(terminal):
+                trace.reset()
+                await trace.run(terminal, None)
+
+
+def _is_swer_terminal(terminal: Terminal) -> bool:
+    return terminal.phases.num_phases == 1
+
+
+def _is_non_swer_terminal(terminal: Terminal) -> bool:
+    return terminal.phases.num_phases > 1
+
+
+def _has_swer_terminal(ce: ConductingEquipment) -> bool:
+    return any(_is_swer_terminal(it) for it in ce.terminals)
+
+
+def _has_non_swer_terminal(ce: ConductingEquipment) -> bool:
+    return any(_is_non_swer_terminal(it) for it in ce.terminals)

zepben/ewb/services/network/tracing/networktrace/__init__.py

@@ -0,0 +1,4 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

zepben/ewb/services/network/tracing/networktrace/actions/__init__.py

@@ -0,0 +1,4 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

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

@@ -0,0 +1,104 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ['EquipmentTreeBuilder']
+
+import uuid
+from typing import Any, Generator
+
+from zepben.ewb.model.cim.iec61970.base.core.conducting_equipment import ConductingEquipment
+from zepben.ewb.services.network.tracing.networktrace.actions.tree_node import TreeNode
+from zepben.ewb.services.network.tracing.networktrace.network_trace_step import NetworkTraceStep
+from zepben.ewb.services.network.tracing.traversal.step_action import StepActionWithContextValue
+from zepben.ewb.services.network.tracing.traversal.step_context import StepContext
+
+EquipmentTreeNode = TreeNode[ConductingEquipment]
+
+
+class EquipmentTreeBuilder(StepActionWithContextValue):
+    """
+    
+    A `StepAction` that can be added to a `NetworkTrace` to build a tree structure representing the paths taken during a trace.
+    The `_roots` are the start items of the trace and the children of a tree node represent the next step paths from a given step in the trace.
+    
+    eg:
+
+    >>> from zepben.ewb import Tracing, NetworkStateOperators
+    >>>
+    >>> tree_builder = EquipmentTreeBuilder()
+    >>> state_operators = NetworkStateOperators.NORMAL
+    >>> (Tracing.network_trace_branching(network_state_operators=state_operators)
+    >>>     .add_condition(state_operators.downstream())
+    >>>     .add_step_action(tree_builder)).run()
+    """
+
+    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())
+
+    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.
+        """
+        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:
+        node = self._roots.get(item.path.to_equipment)
+        if node is None:
+            node = TreeNode(item.path.to_equipment, None)
+            self._roots[item.path.to_equipment] = node
+        return node
+
+    def compute_next_value(
+        self,
+        next_item: NetworkTraceStep[Any],
+        current_item: NetworkTraceStep[Any],
+        current_value: EquipmentTreeNode
+    ) -> EquipmentTreeNode:
+
+        if next_item.path.traced_internally:
+            return current_value
+        else:
+            return TreeNode(next_item.path.to_equipment, current_value)
+
+    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
+
+    def clear(self):
+        self._roots.clear()

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

@@ -0,0 +1,35 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ['TreeNode']
+
+from typing import TypeVar, Generic, Set
+
+T = TypeVar('T')
+
+
+class TreeNode(Generic[T]):
+    """
+    Represents a node in the NetworkTrace tree
+    """
+
+    def __init__(self, identified_object: T, parent=None):
+        self.identified_object = identified_object
+        self._parent: TreeNode = parent
+        self._children: Set[TreeNode] = set()
+
+    @property
+    def parent(self) -> 'TreeNode[T]':
+        return self._parent
+
+    @property
+    def children(self) -> Set['TreeNode[T]']:
+        return set(self._children)
+
+    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)}}}"

zepben/ewb/services/network/tracing/networktrace/compute_data.py

@@ -0,0 +1,60 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ['ComputeData', 'ComputeDataWithPaths']
+
+from typing import TypeVar, Generic
+
+from zepben.ewb.services.network.tracing.networktrace.network_trace_step import NetworkTraceStep
+from zepben.ewb.services.network.tracing.traversal.step_context import StepContext
+
+T = TypeVar('T')
+
+
+class ComputeData(Generic[T]):
+    """
+    Functional interface used to compute contextual data stored on a NetworkTraceStep.
+    """
+
+    def __init__(self, func):
+        self._func = func
+
+    def compute_next(self, current_step: NetworkTraceStep[T], current_context: StepContext, next_path: NetworkTraceStep.Path) -> T:
+        """
+        Called for each new NetworkTraceStep in a NetworkTrace. The value returned from this function
+        will be stored against the next step within NetworkTraceStep. data.
+
+        `currentStep` - The current step of the trace.
+        `currentContext` - The context of teh current step in the trace.
+        `nextPath` - The next path of the next NetworkTraceStep that the data will be associated with.
+
+        Returns The data to associate with the next NetworkTraceStep.
+        """
+        return self._func(current_step, current_context, next_path)
+
+
+class ComputeDataWithPaths(Generic[T]):
+    """
+    Functional interface used to compute contextual data stored on a NetworkTraceStep. This can be used when the
+    contextual data can only be computed by knowing all the next paths that can be stepped to from a given step.
+    """
+
+    def __init__(self, func):
+        self._func = func or (lambda *args: None)
+
+    def compute_next(self, current_step: NetworkTraceStep[T], current_context: StepContext, next_path: NetworkTraceStep.Path,
+                     next_paths: list[NetworkTraceStep.Path]) -> T:
+        """
+        Called for each new NetworkTraceStep in a NetworkTrace. The value returned from this function
+        will be stored against the next step within NetworkTraceStep. data.
+
+        `currentStep` - The current step of the trace.
+        `currentContext` - The context of teh current step in the trace.
+        `nextPath` - The next path of the next NetworkTraceStep that the data will be associated with.
+        `nextPaths` - A list of all the next paths that the current step can trace to.
+
+        Returns The data to associate with the next NetworkTraceStep.
+        """
+        return self._func(current_step, current_context, next_path, next_paths)

zepben/ewb/services/network/tracing/networktrace/conditions/__init__.py

@@ -0,0 +1,4 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

zepben/ewb/services/network/tracing/networktrace/conditions/conditions.py

@@ -0,0 +1,73 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+from __future__ import annotations
+
+__all__ = ['upstream', 'downstream', 'with_direction', 'limit_equipment_steps', 'stop_at_open']
+
+from typing import TYPE_CHECKING, TypeVar, Type, Callable
+
+from zepben.ewb.services.network.tracing.feeder.feeder_direction import FeederDirection
+from zepben.ewb.services.network.tracing.networktrace.conditions.direction_condition import DirectionCondition
+from zepben.ewb.services.network.tracing.networktrace.conditions.equipment_step_limit_condition import EquipmentStepLimitCondition
+from zepben.ewb.services.network.tracing.networktrace.conditions.equipment_type_step_limit_condition import EquipmentTypeStepLimitCondition
+
+T = TypeVar('T')
+
+if TYPE_CHECKING:
+    from zepben.ewb.services.network.tracing.networktrace.network_trace_step import NetworkTraceStep
+    from zepben.ewb.services.network.tracing.traversal.queue_condition import QueueCondition
+    from zepben.ewb.services.network.tracing.traversal.stop_condition import StopCondition
+    from zepben.ewb import ConductingEquipment, NetworkStateOperators
+
+    DSLLambda = Callable[[NetworkStateOperators], QueueCondition[NetworkTraceStep[T]]]
+
+
+def upstream() -> DSLLambda:
+    """
+    Creates a [NetworkTrace] condition that will cause tracing a feeder upstream (towards the head terminal).
+    This uses [FeederDirectionStateOperations.get_direction] receiver instance method within the condition.
+
+    :return: [NetworkTraceQueueCondition] that results in upstream tracing.
+    """
+    return lambda state_operator: state_operator.with_direction(FeederDirection.UPSTREAM)
+
+
+def downstream() -> DSLLambda:
+    """
+    Creates a [NetworkTrace] condition that will cause tracing a feeder downstream (away from the head terminal).
+    This uses [FeederDirectionStateOperations.get_direction] receiver instance method within the condition.
+
+    :return: [NetworkTraceQueueCondition] that results in downstream tracing.
+    """
+    return lambda state_operator: state_operator.with_direction(FeederDirection.DOWNSTREAM)
+
+
+def with_direction(direction: FeederDirection) -> DSLLambda:
+    """
+    Creates a [NetworkTrace] condition that will cause tracing only terminals with directions that match [direction].
+    This uses [FeederDirectionStateOperations.get_direction] receiver instance method within the condition.
+
+    :return: [NetworkTraceQueueCondition] that results in upstream tracing.
+    """
+    return lambda state_operator: DirectionCondition(direction, state_operator)
+
+
+def limit_equipment_steps(limit: int, equipment_type: Type[ConductingEquipment] = None) -> StopCondition[NetworkTraceStep[T]]:
+    """
+    Creates a [NetworkTrace] condition that stops tracing a path once a specified number of equipment steps have been reached.
+
+    :param limit: The maximum number of equipment steps allowed before stopping.
+    :param equipment_type: The class of the equipment type to track against the limit
+
+    :return: A [NetworkTraceStopCondition] that stops tracing the path once the step limit is reached.
+    """
+    if equipment_type is not None:
+        return EquipmentTypeStepLimitCondition(limit, equipment_type)
+    return EquipmentStepLimitCondition(limit)
+
+
+def stop_at_open():
+    return lambda state_operator: state_operator.stop_at_open()

zepben/ewb/services/network/tracing/networktrace/conditions/direction_condition.py

@@ -0,0 +1,63 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+from __future__ import annotations
+
+__all__ = ['DirectionCondition']
+
+from typing import TypeVar, TYPE_CHECKING, Generic, Type
+
+from zepben.ewb import require
+from zepben.ewb.model.cim.iec61970.base.wires.clamp import Clamp
+from zepben.ewb.model.cim.iec61970.base.wires.cut import Cut
+from zepben.ewb.services.network.tracing.feeder.feeder_direction import FeederDirection
+from zepben.ewb.services.network.tracing.networktrace.network_trace_step import NetworkTraceStep
+from zepben.ewb.services.network.tracing.traversal.queue_condition import QueueCondition
+
+if TYPE_CHECKING:
+    from zepben.ewb.services.network.tracing.networktrace.operators.network_state_operators import NetworkStateOperators
+    from zepben.ewb import StepContext
+
+T = TypeVar('T')
+
+
+class DirectionCondition(QueueCondition[NetworkTraceStep[T]], Generic[T]):
+
+    def __init__(self, direction: FeederDirection, state_operators: Type[NetworkStateOperators]):
+        require(direction != FeederDirection.CONNECTOR, lambda: 'A direction of CONNECTOR is not currently supported')
+        self.direction = direction
+        self.state_operators = state_operators
+        self.get_direction = self.state_operators.get_direction
+
+    def should_queue(self, next_item: NetworkTraceStep[T], next_context: StepContext[T], current_item: NetworkTraceStep[T],
+                     current_context: StepContext[T]) -> bool:
+        return self._should_queue(next_item.path)
+
+    def _should_queue(self, path: NetworkTraceStep.Path) -> bool:
+        # Cuts do weird things with directions depending on if they are energised from an external connection, or through a "closed" cut. To prevent
+        # dealing with this awful mess, it is much simpler to just ask if anything else past it needs queueing. This could be made to short-circuit
+        # for traversing downstream, but the code is much more complex to only save one extra step.
+        if isinstance(path.to_equipment, Cut):
+            return self._should_queue_next_paths(path)
+        elif path.traced_internally or path.did_traverse_ac_line_segment:
+            return self.direction in self.get_direction(path.to_terminal)
+        else:
+            return self.direction.complementary_external_direction in self.get_direction(path.to_terminal)
+
+    def should_queue_start_item(self, item: NetworkTraceStep[T]) -> bool:
+        if self.direction in self.get_direction(item.path.to_terminal):
+            return True
+        # Because cuts and clamps behave a bit different with directions than other equipment terminals, we can also check if any further paths needs to be
+        # queued, and if they do we queue the start item.
+        elif isinstance(item.path.to_equipment, (Clamp, Cut)):
+            return self._should_queue_next_paths(item.path)
+        return False
+
+    def _should_queue_next_paths(self, path: NetworkTraceStep.Path) -> bool:
+        for next_path in self.state_operators.next_paths(path):
+            if not (next_path.traced_internally and self.state_operators.is_open(path.to_equipment)):
+                if self._should_queue(next_path):
+                    return True
+        return False

zepben/ewb/services/network/tracing/networktrace/conditions/equipment_step_limit_condition.py

@@ -0,0 +1,26 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+from __future__ import annotations
+
+__all__ = ['EquipmentStepLimitCondition']
+
+from typing import TypeVar, TYPE_CHECKING
+
+from zepben.ewb.services.network.tracing.traversal.stop_condition import StopCondition
+
+T = TypeVar('T')
+
+if TYPE_CHECKING:
+    from zepben.ewb import StepContext, NetworkTraceStep
+
+
+class EquipmentStepLimitCondition(StopCondition['NetworkTraceStep[T]']):
+    def __init__(self, limit: int):
+        super().__init__()
+        self.limit = limit
+
+    def should_stop(self, item: NetworkTraceStep[T], context: StepContext) -> bool:
+        return item.num_equipment_steps >= self.limit

zepben/ewb/services/network/tracing/networktrace/conditions/equipment_type_step_limit_condition.py

@@ -0,0 +1,44 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+from __future__ import annotations
+
+__all__ = ['EquipmentTypeStepLimitCondition']
+
+from typing import Generic, TypeVar, TYPE_CHECKING, Type
+
+from zepben.ewb.services.network.tracing.traversal.context_value_computer import ContextValueComputer
+from zepben.ewb.services.network.tracing.traversal.stop_condition import StopConditionWithContextValue
+
+if TYPE_CHECKING:
+    from zepben.ewb import ConductingEquipment, StepContext, NetworkTraceStep
+
+T = TypeVar('T')
+U = TypeVar('U')
+
+
+class EquipmentTypeStepLimitCondition(StopConditionWithContextValue[T], Generic[T, U]):
+    def __init__(self, limit: int, equipment_type: Type[ConductingEquipment]):
+        StopConditionWithContextValue.__init__(self, self.should_stop)
+        ContextValueComputer.__init__(self, f'sdk:{equipment_type.name}Count')
+        self.limit = limit
+        self.equipment_type = equipment_type
+
+    def should_stop(self, item: NetworkTraceStep[T], context: StepContext) -> bool:
+        return self.get_context_value(context) >= self.limit
+
+    def compute_initial_value(self, item: NetworkTraceStep[T]) -> int:
+        return 0
+
+    def compute_next_value(self, next_item: NetworkTraceStep[T], current_item: NetworkTraceStep[T], current_value: int) -> int:
+        if next_item.path.traced_internally:
+            return current_value
+        if self.matches_equipment_type(next_item.path.to_equipment):
+            return current_value + 1
+        else:
+            return current_value
+
+    def matches_equipment_type(self, conducting_equipment: ConductingEquipment) -> bool:
+        return isinstance(conducting_equipment, self.equipment_type)

zepben/ewb/services/network/tracing/networktrace/conditions/network_trace_queue_condition.py

@@ -0,0 +1,67 @@
+#  Copyright 2025 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ['NetworkTraceQueueCondition']
+
+from typing import TypeVar, Generic
+
+from zepben.ewb.services.network.tracing.networktrace.network_trace_step import NetworkTraceStep
+from zepben.ewb.services.network.tracing.traversal.queue_condition import QueueCondition, ShouldQueue
+from zepben.ewb.services.network.tracing.traversal.step_context import StepContext
+
+T = TypeVar('T')
+
+
+class NetworkTraceQueueCondition(QueueCondition[NetworkTraceStep[T]], Generic[T]):
+    """
+    A special queue condition implementation that allows only checking `should_queue` when a [NetworkTraceStep] matches a given
+    [NetworkTraceStep.Type]. When [step_type] is:
+    *[NetworkTraceStep.Type.ALL]: [should_queue] will be called for every step.
+    *[NetworkTraceStep.Type.INTERNAL]: [shouldQueue] will be called only when [NetworkTraceStep.type] is [NetworkTraceStep.Type.INTERNAL].
+    *[NetworkTraceStep.Type.EXTERNAL]: [shouldQueue] will be called only when [NetworkTraceStep.type] is [NetworkTraceStep.Type.EXTERNAL].
+
+    If the step does not match the given step type, `true` will always be returned.
+    """
+
+    def __init__(self, step_type: NetworkTraceStep.Type, condition: ShouldQueue = None):
+        """
+        :param step_type: The step type to match to check `should_queue`.
+        :param condition: function with the signature of `ShouldQueue` to be called when step_type matches the current items step
+        """
+        super().__init__(self.should_queue)
+        if condition is not None:
+            self.should_queue_matched_step = condition
+        self.should_queue = self._should_queue_func(step_type)
+
+    def should_queue(self, next_item: T, next_context: StepContext, current_item: T, current_context: StepContext) -> bool:
+        raise NotImplementedError()
+
+    def should_queue_matched_step(self, next_item: NetworkTraceStep[T], next_context: StepContext, current_item: NetworkTraceStep[T],
+                                  current_context: StepContext) -> bool:
+        """
+        The logic you would normally put in `should_queue`. However, this will only be called when a step matches the `step_type`
+        """
+        raise NotImplementedError()
+
+    def should_queue_internal_step(self, next_item: NetworkTraceStep[T], next_context: StepContext, current_item: NetworkTraceStep[T],
+                                   current_context: StepContext) -> bool:
+        if next_item.type() == NetworkTraceStep.Type.INTERNAL:
+            return self.should_queue_matched_step(next_item, next_context, current_item, current_context)
+        return True
+
+    def should_queue_external_step(self, next_item: NetworkTraceStep[T], next_context: StepContext, current_item: NetworkTraceStep[T],
+                                   current_context: StepContext) -> bool:
+        if next_item.type() == NetworkTraceStep.Type.EXTERNAL:
+            return self.should_queue_matched_step(next_item, next_context, current_item, current_context)
+        return True
+
+    def _should_queue_func(self, step_type: NetworkTraceStep.Type) -> ShouldQueue:
+        if step_type == NetworkTraceStep.Type.ALL:
+            return self.should_queue_matched_step
+        elif step_type == NetworkTraceStep.Type.INTERNAL:
+            return self.should_queue_internal_step
+        elif step_type == NetworkTraceStep.Type.EXTERNAL:
+            return self.should_queue_external_step
+        raise ValueError(f"INTERNAL ERROR: step type [{step_type}] didn't match expected")