Fix slow traces

changelog.md

@@ -5,10 +5,10 @@
 * `AcLineSegment` supports adding a maximum of 2 terminals. Mid-span terminals are no longer supported and models should migrate to using `Clamp`.
 * `Clamp` supports only adding a single terminal.
 * `FeederDirectionStateOperations` have been reworked to take `NetworkStateOperators` as a parameter.
-* `RemoveDirection` has been removed. It did not work reliably with dual fed networks with loops. You now need to clear direction using the new
-`ClearDirection` and reapply directions where appropriate using `SetDirection`.
+* `RemoveDirection` has been removed. It did not work reliably with dual fed networks with loops. You now need to clear direction using the new `ClearDirection`
+  and reapply directions where appropriate using `SetDirection`.
 * `Cut` supports adding a maximum of 2 terminals.
-
+* `NetworkTraceTracker` now uses a `set` to track visited objects, if you were using unhashable objects this will need to be addressed.
 
 ### New Features
 * Added `ClearDirection` that clears feeder directions.
@@ -22,9 +22,8 @@
   * You can now add sites to the `TestNetworkBuilder` via `addSite`.
   * You can now add busbar sections natively with `from_busbar_section` and `to_busbar_section`
   * The prefix for generated mRIDs for "other" equipment can be specified with the `default_mrid_prefix` argument in `from_other` and `to_other`.
-* When processing feeder assignments, all LV feeders belonging to a dist substation site will now be considered energized when the site is energized by a feeder.
-
-
+* When processing feeder assignments, all LV feeders belonging to a dist substation site will now be considered energized when the site is energized by a
+  feeder.
 
 ### Fixes
 * When finding `LvFeeders` in the `Site` we will now exclude `LvFeeders` that start with an open `Switch`
@@ -40,8 +39,6 @@
 * `AssignToFeeders`/`AssignToLvFeeders` now finds back-fed equipment correctly
 * `AssignToFeeders` and `AssignToLvFeeders` will now associate `PowerElectronicUnits` with their `powerElectronicsConnection` `Feeder`/`LvFeeder`.
 
-
-
 ### Notes
 * None.
 
@@ -95,7 +92,8 @@
 
 ## [0.44.0] - 2025-01-24
 ### Breaking Changes
-* `GrpcChannelBuilder.build()` now accepts a `timeout_seconds` argument. This is the timeout used for each connection attempt so the total amount of time the connection test may take to fail can be greater than `timeout_seconds`.
+* `GrpcChannelBuilder.build()` now accepts a `timeout_seconds` argument. This is the timeout used for each connection attempt so the total amount of time the
+  connection test may take to fail can be greater than `timeout_seconds`.
 
 ### New Features
 * Added the following new CIM classes:
@@ -131,9 +129,10 @@
 * Removed `ProcessingPaused` current state response message as this functionality won't be supported.
 * `QueryNetworkStateClient.get_current_states` now returns a `CurrentStateEventBatch` rather than just the events themselves.
 * `QueryNetworkStateService.on_get_current_states` must now return a stream of `CurrentStateEventBatch` rather than just the events themselves.
-* `AcLineSegment.per_length_sequence_impedance` has been corrected to `per_length_impedance`. This has been done in a non-breaking way, however the public 
+* `AcLineSegment.per_length_sequence_impedance` has been corrected to `per_length_impedance`. This has been done in a non-breaking way, however the public
   resolver `Resolvers.per_length_sequence_impedance` is now `Resolvers.per_length_impedance`, correctly reflecting the CIM relationship.
-* Removed `get_current_equipment_for_feeder` implementation for `NetworkConsumerClient` as its functionality is now incorporated in `get_equipment_for_container`.
+* Removed `get_current_equipment_for_feeder` implementation for `NetworkConsumerClient` as its functionality is now incorporated in
+  `get_equipment_for_container`.
 
 ### New Features
 * Added `BatchNotProcessed` current state response. This is used to indicate a batch has been ignored, rather than just returning a `BatchSuccessful`.

src/zepben/evolve/services/network/tracing/networktrace/network_trace.py

@@ -5,7 +5,7 @@
 
 from collections.abc import Callable
 from functools import singledispatchmethod
-from typing import TypeVar, Union, Generic, Set, Type, Generator
+from typing import TypeVar, Union, Generic, Set, Type, Generator, FrozenSet
 
 from zepben.evolve.model.cim.iec61970.base.wires.clamp import Clamp
 from zepben.evolve.model.cim.iec61970.base.wires.aclinesegment import AcLineSegment
@@ -295,15 +295,15 @@ def create_new_this(self) -> 'NetworkTrace[T]':
     def start_nominal_phase_path(phases: PhaseCode) -> Set[NominalPhasePath]:
         return {NominalPhasePath(it, it) for it in phases.single_phases} if phases and phases.single_phases else set()
 
-    def has_visited(self, terminal: Terminal, phases: set[SinglePhaseKind]) -> bool:
+    def has_visited(self, terminal: Terminal, phases: FrozenSet[SinglePhaseKind]) -> bool:
         parent = self.parent
         while parent is not None:
             if parent._tracker.has_visited(terminal, phases):
                 return True
             parent = parent.parent
         return self._tracker.has_visited(terminal, phases)
 
-    def visit(self, terminal: Terminal, phases: set[SinglePhaseKind]) -> bool:
+    def visit(self, terminal: Terminal, phases: FrozenSet[SinglePhaseKind]) -> bool:
         if self.parent and self.parent.has_visited(terminal, phases):
                 return False
         return self._tracker.visit(terminal, phases)

src/zepben/evolve/services/network/tracing/networktrace/network_trace_step.py

@@ -6,7 +6,7 @@
 
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Set, Generic, TypeVar, TYPE_CHECKING, Optional, List
+from typing import Set, Generic, TypeVar, TYPE_CHECKING, Optional, List, FrozenSet
 
 from zepben.evolve.services.network.tracing.connectivity.nominal_phase_path import NominalPhasePath
 
@@ -50,10 +50,10 @@ class Path:
         traversed_ac_line_segment: Optional[AcLineSegment] = field(default=None)
         nominal_phase_paths: Optional[Set[NominalPhasePath]] = field(default_factory=set)
 
-        def to_phases_set(self) -> Set[SinglePhaseKind]:
+        def to_phases_set(self) -> FrozenSet[SinglePhaseKind]:
             if len(self.nominal_phase_paths) == 0:
-                return set()
-            return set(map(lambda it: it.to_phase, self.nominal_phase_paths))
+                return frozenset()
+            return frozenset(map(lambda it: it.to_phase, self.nominal_phase_paths))
 
         @property
         def from_equipment(self) -> ConductingEquipment:

src/zepben/evolve/services/network/tracing/networktrace/network_trace_tracker.py

@@ -2,7 +2,7 @@
 #  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 typing import Set, Any
+from typing import Set, Any, FrozenSet
 
 from zepben.evolve.model.cim.iec61970.base.core.terminal import Terminal
 from zepben.evolve.model.cim.iec61970.base.wires.single_phase_kind import SinglePhaseKind
@@ -13,17 +13,17 @@ class NetworkTraceTracker:
     Internal class that tracks visited state of a Terminal's Phase in a Network Trace
     """
     def __init__(self):
-        self._visited = list()
+        self._visited = set()
 
-    def has_visited(self, terminal: Terminal, phases: Set[SinglePhaseKind]=None) -> bool:
+    def has_visited(self, terminal: Terminal, phases: FrozenSet[SinglePhaseKind]=None) -> bool:
         """Returns True if this Terminal's Phase has been visited, False otherwise"""
         return self._get_key(terminal, phases) in self._visited
 
-    def visit(self, terminal: Terminal, phases: Set[SinglePhaseKind]=None) -> bool:
+    def visit(self, terminal: Terminal, phases: FrozenSet[SinglePhaseKind]=None) -> bool:
         """Marks this Terminal's Phase as visited"""
         key = self._get_key(terminal, phases)
         if key not in self._visited:
-            self._visited.append(self._get_key(terminal, phases))
+            self._visited.add(self._get_key(terminal, phases))
             return True
         return False
 
@@ -32,8 +32,8 @@ def clear(self):
         self._visited.clear()
 
     @staticmethod
-    def _get_key(terminal: Terminal, phases: Set[SinglePhaseKind]) -> Any:
+    def _get_key(terminal: Terminal, phases: FrozenSet[SinglePhaseKind]) -> Any:
         if phases:
-            return terminal, phases
+            return terminal.mrid, phases
         else:
-            return terminal
+            return terminal.mrid

test/services/network/tracing/networktrace/test_network_trace.py

@@ -190,7 +190,7 @@ async def test_can_run_large_branching_traces(self):
                 builder.from_junction(num_terminals=1) \
                        .to_acls()
 
-                for i in range(500):
+                for i in range(1000):
                     builder.to_junction(mrid=f'junc-{i}', num_terminals=3) \
                            .to_acls(mrid=f'acls-{i}-top') \
                            .from_acls(mrid=f'acls-{i}-bottom') \