diff --git a/src/zepben/ewb/services/network/tracing/networktrace/network_trace.py b/src/zepben/ewb/services/network/tracing/networktrace/network_trace.py
index 070b607a9..c5c888184 100644
--- a/src/zepben/ewb/services/network/tracing/networktrace/network_trace.py
+++ b/src/zepben/ewb/services/network/tracing/networktrace/network_trace.py
@@ -413,9 +413,7 @@ def visit(self, terminal: Terminal, phases: FrozenSet[SinglePhaseKind]) -> bool:
         return self._tracker.visit(terminal, phases)
 
 
-def default_condition_step_type(step_type: CanActionItem):
-    if step_type is None:
-        return False
+def default_condition_step_type(step_type: CanActionItem) -> NetworkTraceStep.Type:
     if step_type == NetworkTraceActionType.ALL_STEPS:
         return NetworkTraceStep.Type.ALL
     elif step_type == NetworkTraceActionType.FIRST_STEP_ON_EQUIPMENT:
@@ -432,19 +430,4 @@ def compute_data_with_action_type(compute_data: ComputeData[T], action_type: Can
                 current_step.data if next_path.traced_internally else compute_data.compute_next(current_step, current_context, next_path)
             )
         )
-    raise Exception(f'{action_type.__class__}: step doesnt match expected types')
-
-
-def with_paths_with_action_type(self, action_type: NetworkTraceActionType) -> ComputeDataWithPaths[T]:
-    if action_type == NetworkTraceActionType.ALL_STEPS:
-        return self
-    elif action_type == NetworkTraceActionType.FIRST_STEP_ON_EQUIPMENT:
-        return ComputeDataWithPaths(
-            lambda current_step, current_context, next_path, next_paths: (
-                current_step.data if next_path.traced_internally else self.compute_next(current_step, current_context, next_path, next_paths)
-            )
-        )
-    raise Exception('step doesnt match expected types')
-
-
-ComputeDataWithPaths[T].with_action_type = with_paths_with_action_type
+    raise Exception(f'{action_type.__name__}: step doesnt match expected types')
diff --git a/src/zepben/ewb/services/network/tracing/networktrace/network_trace_action_type.py b/src/zepben/ewb/services/network/tracing/networktrace/network_trace_action_type.py
index f672393f8..a03728aac 100644
--- a/src/zepben/ewb/services/network/tracing/networktrace/network_trace_action_type.py
+++ b/src/zepben/ewb/services/network/tracing/networktrace/network_trace_action_type.py
@@ -2,14 +2,14 @@
 #  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 Callable, Set, Any
+from typing import Callable, FrozenSet, Any
 from enum import Enum
 
 from zepben.ewb import Terminal, SinglePhaseKind
 from zepben.ewb.services.network.tracing.networktrace.network_trace_step import NetworkTraceStep
 from zepben.ewb.services.network.tracing.traversal.step_context import StepContext
 
-HasTracked = Callable[[Terminal, Set[SinglePhaseKind]], bool]
+HasTracked = Callable[[Terminal, FrozenSet[SinglePhaseKind]], bool]
 CanActionItem = Callable[[NetworkTraceStep[Any], StepContext, HasTracked], bool]
 
 
diff --git a/src/zepben/ewb/services/network/tracing/networktrace/network_trace_step.py b/src/zepben/ewb/services/network/tracing/networktrace/network_trace_step.py
index c72ceacc3..bcaab38b9 100644
--- a/src/zepben/ewb/services/network/tracing/networktrace/network_trace_step.py
+++ b/src/zepben/ewb/services/network/tracing/networktrace/network_trace_step.py
@@ -95,7 +95,10 @@ def did_traverse_ac_line_segment(self) -> bool:
         def next_num_equipment_steps(self, current_num: int) -> int:
             return current_num + 1 if self.traced_externally else current_num
 
-    Type = Enum('Type', ('ALL', 'INTERNAL', 'EXTERNAL'))
+    class Type(Enum):
+        ALL = 'ALL'
+        INTERNAL = 'INTERNAL'
+        EXTERNAL = 'EXTERNAL'
 
     def __init__(self, path: Path, num_terminal_steps: int, num_equipment_steps: int, data: T):
         self.path = path