PR Changes

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

Author: chestm007

src/zepben/evolve/services/network/tracing/phases/set_phases.py

@@ -90,7 +90,7 @@ def _terminals_from_network():
     async def _(
         self,
         start_terminal: Terminal,
-        phases: Union[PhaseCode, List[SinglePhaseKind]]=None,
+        phases: Union[PhaseCode, List[SinglePhaseKind], Set[SinglePhaseKind]]=None,
         network_state_operators: Type[NetworkStateOperators]=NetworkStateOperators.NORMAL,
         seed_terminal: Terminal=None):
         """
@@ -144,11 +144,8 @@ def spread_phases(
         :param network_state_operators: The `NetworkStateOperators` to be used when setting phases.
         """
 
-        if phases is None:
-            self.spread_phases(from_terminal, to_terminal, from_terminal.phases.single_phases, network_state_operators)
-        else:
-            paths = self._get_nominal_phase_paths(network_state_operators, from_terminal, to_terminal, phases)
-            self._flow_phases(network_state_operators, from_terminal, to_terminal, paths)
+        paths = self._get_nominal_phase_paths(network_state_operators, from_terminal, to_terminal, phases or from_terminal.phases.single_phases)
+        self._flow_phases(network_state_operators, from_terminal, to_terminal, paths)
 
     async def _run_terminals(self, terminals: Iterable[Terminal], network_state_operators: Type[NetworkStateOperators]):
 
@@ -356,14 +353,15 @@ def _flow_transformer_phases(
         # Split the phases into ones we need to flow directly, and ones that have been added by a transformer. In
         #  the case of an added Y phase (SWER -> LV2 transformer) we need to flow the phases before we can calculate
         #  the missing phase.
+        flow_phases = (p for p in paths if p.from_phase == SinglePhaseKind.NONE)
+        add_phases = (p for p in paths if p.from_phase != SinglePhaseKind.NONE)
+        for p in flow_phases:
+            self._try_add_phase(from_terminal, from_phases, to_terminal, to_phases, p.to_phase, allow_suspect_flow,
+                                lambda: updated_phases.append(True))
 
-        for path in paths:
-            if path.from_phase == SinglePhaseKind.NONE:
-                self._try_add_phase(from_terminal, from_phases, to_terminal, to_phases, path.to_phase, allow_suspect_flow,
-                                    lambda: updated_phases.append(True))
-            else:
-                self._try_set_phase(from_phases[path.from_phase], from_terminal, from_phases, path.from_phase,
-                                    to_terminal, to_phases, path.to_phase, lambda: updated_phases.append(True))
+        for p in add_phases:
+            self._try_set_phase(from_phases[p.from_phase], from_terminal, from_phases, p.from_phase,
+                                to_terminal, to_phases, p.to_phase, lambda: updated_phases.append(True))
 
         return any(updated_phases)
 
@@ -399,8 +397,7 @@ def _try_set_phase(
         on_success: Callable[[], None]):
 
         try:
-            if phase != SinglePhaseKind.NONE:
-                to_phases[to_] = phase
+            if phase != SinglePhaseKind.NONE and to_phases.__setitem__(to_, phase):
                 if self._debug_logger:
                     self._debug_logger.info(f'   {from_terminal.mrid}[{from_}] -> {to_terminal.mrid}[{to_}]: set to {phase}')
                 on_success()

src/zepben/evolve/services/network/tracing/traversal/debug_logging.py

@@ -85,18 +85,15 @@ def _wrap_attr(_index: int, _attr: str, _msg: str) -> None:
             :param _attr: Method/Function name.
             :raises AttributeError: if ``wrappable`` is already wrapped
             """
+
             if isinstance(_attr, tuple):
                 _attr_name, _log_attr_name = _attr
             else:
                 _attr_name = _log_attr_name = _attr
 
-            # Wrapped classes will have __wrapped__ == True - if it exists on the obj passed in, the user is attempting to wrap an
-            # already wrapped object. This can lead to unexpected outcomes so we do not support it
-            if (to_wrap := getattr(w_obj, _attr_name)) and hasattr(to_wrap, '__wrapped__'):
-                    raise AttributeError(f'Wrapped objects cannot be rewrapped, pass in the original object instead.')
+            to_wrap = getattr(w_obj, _attr_name)
 
             setattr(w_obj, _attr_name, self._log_method_call(to_wrap, f'{self.description}: {_log_attr_name}({_index})' + _msg))
-            setattr(w_obj, '__wrapped__', True)
 
         for clazz in (StepAction, StopCondition, QueueCondition):
             if isinstance(w_obj, clazz):

src/zepben/evolve/services/network/tracing/traversal/step_action.py

@@ -25,7 +25,7 @@ class StepAction(Generic[T]):
     """
 
     def __init__(self, _func: StepActionFunc = None):
-        self._func = _func
+        self._func = _func or self._apply
 
     def __init_subclass__(cls):
         """
@@ -50,6 +50,17 @@ def apply(self, item: T, context: StepContext):
 
         return self._func(item, context)
 
+    @abstractmethod
+    def _apply(self, item: T, context: StepContext):
+        """
+        Override this method instead of ``self.apply`` directly
+
+        :param item: The current item in the traversal.
+        :param context: The context associated with the current traversal step.
+        """
+        raise NotImplementedError()
+
+
 class StepActionWithContextValue(StepAction[T], ContextValueComputer[T]):
     """
     Interface representing a step action that utilises a value stored in the :class:`StepContext`.
@@ -59,19 +70,9 @@ class StepActionWithContextValue(StepAction[T], ContextValueComputer[T]):
     """
 
     def __init__(self, key: str, _func: StepActionFunc = None):
-        StepAction.__init__(self, _func or self._apply)
+        StepAction.__init__(self, _func)
         ContextValueComputer.__init__(self, key)
 
-    @abstractmethod
-    def _apply(self, item: T, context: StepContext):
-        """
-        Override this method instead of ``self.apply`` directly
-
-        :param item: The current item in the traversal.
-        :param context: The context associated with the current traversal step.
-        """
-        raise NotImplementedError()
-
     @abstractmethod
     def compute_initial_value(self, item: T):
         raise NotImplementedError()

src/zepben/evolve/services/network/tracing/traversal/traversal.py

@@ -26,7 +26,6 @@
 __all__ = ["Traversal"]
 
 from zepben.evolve.services.network.tracing.traversal.queue import TraversalQueue
-from zepben.evolve.util import extra_kwargs_not_allowed
 
 T = TypeVar('T')
 U = TypeVar('U')
@@ -205,21 +204,20 @@ def create_new_this(self) -> D:
         raise NotImplementedError
 
     @singledispatchmethod
-    def add_condition(self, condition: ConditionTypes, **kwargs) -> D:
+    def add_condition(self, condition: ConditionTypes) -> D:
         """
         Adds a traversal condition to the traversal.
 
         :param condition: The condition to add.
-        :keyword allow_re_wrapping: Allow rewrapping of :class:`StopConditions` with debug logging
 
         :return: this traversal instance.
         """
 
         if callable(condition):  # Callable[[NetworkTraceStep[T], StepContext], None]
             if len(inspect.getfullargspec(condition).args) == 2:
-                return self.add_stop_condition(condition, **kwargs)
+                return self.add_stop_condition(condition)
             elif len(inspect.getfullargspec(condition).args) == 4:
-                return self.add_queue_condition(condition, **kwargs)
+                return self.add_queue_condition(condition)
             else:
                 raise RuntimeError(f'Condition does not match expected: Number of args is not 2(Stop Condition) or 4(QueueCondition)')
 
@@ -231,30 +229,27 @@ def add_condition(self, condition: ConditionTypes, **kwargs) -> D:
 
     @singledispatchmethod
     @add_condition.register(StopCondition)
-    def add_stop_condition(self, condition: StopConditionTypes, **kwargs) -> D:
+    def add_stop_condition(self, condition: StopConditionTypes) -> D:
         """
         Adds a stop condition to the traversal. If any stop condition returns
         ``True``, the traversal will not call the callback to queue more items
         from the current item.
 
         :param condition: The stop condition to add.
-        :keyword allow_re_wrapping: Allow rewrapping of :class:`StopCondition`s with debug logging
         :return: this traversal instance.
         """
 
         raise RuntimeError(f'Condition [{condition.__class__.__name__}] does not match expected: [StopCondition | StopConditionWithContextValue | Callable]')
 
     @add_stop_condition.register(Callable)
-    def _(self, condition: ShouldStop, **kwargs):
-        return self.add_stop_condition(StopCondition(condition), **kwargs)
+    def _(self, condition: ShouldStop):
+        return self.add_stop_condition(StopCondition(condition))
 
     @add_stop_condition.register
-    def _(self, condition: StopCondition, **kwargs):
+    def _(self, condition: StopCondition):
 
         if self._debug_logger is not None:
-            self._debug_logger.wrap(condition, kwargs.pop('allow_re_wrapping', False))
-
-        extra_kwargs_not_allowed(kwargs, 'add_stop_condition')
+            self._debug_logger.wrap(condition)
 
         self.stop_conditions.append(condition)
         if isinstance(condition, StopConditionWithContextValue):
@@ -281,30 +276,27 @@ def matches_any_stop_condition(self, item: T, context: StepContext) -> bool:
 
     @add_condition.register(QueueCondition)
     @singledispatchmethod
-    def add_queue_condition(self, condition: QueueConditionTypes, **kwargs) -> D:
+    def add_queue_condition(self, condition: QueueConditionTypes) -> D:
         """
         Adds a queue condition to the traversal.
         Queue conditions determine whether an item should be queued for traversal.
         All registered queue conditions must return true for an item to be queued.
 
         :param condition: The queue condition to add.
-        :keyword allow_re_wrapping: Allow rewrapping of :class:`QueueCondition`s with debug logging
         :returns: The current traversal instance.
         """
 
         raise RuntimeError(f'Condition [{condition.__class__.__name__}] does not match expected: [QueueCondition | QueueConditionWithContextValue | Callable]')
 
     @add_queue_condition.register(Callable)
-    def _(self, condition: ShouldQueue, **kwargs):
-        return self.add_queue_condition(QueueCondition(condition), **kwargs)
+    def _(self, condition: ShouldQueue):
+        return self.add_queue_condition(QueueCondition(condition))
 
     @add_queue_condition.register
-    def _(self, condition: QueueCondition, **kwargs):
+    def _(self, condition: QueueCondition):
 
         if self._debug_logger is not None:
-            self._debug_logger.wrap(condition, kwargs.pop('allow_re_wrapping', False))
-
-        extra_kwargs_not_allowed(kwargs, 'add_queue_condition')
+            self._debug_logger.wrap(condition)
 
         self.queue_conditions.append(condition)
         if isinstance(condition, QueueConditionWithContextValue):
@@ -324,24 +316,21 @@ def copy_queue_conditions(self, other: Traversal[T, D]) -> D:
         return self
 
     @singledispatchmethod
-    def add_step_action(self, action: StepActionTypes, **kwargs) -> D:
+    def add_step_action(self, action: StepActionTypes) -> D:
         """
         Adds an action to be performed on each item in the traversal, including the
         starting items.
 
         :param action: The action to perform on each item.
-        :keyword allow_re_wrapping: Allow rewrapping of :class:`StepAction`s with debug logging
         :return: The current traversal instance.
         """
 
         raise RuntimeError(f'StepAction [{action.__class__.__name__}] does not match expected: [StepAction | StepActionWithContextValue | Callable]')
 
     @add_step_action.register
-    def _(self, action: StepAction, **kwargs):
+    def _(self, action: StepAction):
         if self._debug_logger is not None:
-            self._debug_logger.wrap(action, kwargs.pop('allow_re_wrapping', False))
-
-        extra_kwargs_not_allowed(kwargs, 'add_step_action')
+            self._debug_logger.wrap(action)
 
         self.step_actions.append(action)
         if isinstance(action, StepActionWithContextValue):
@@ -350,47 +339,45 @@ def _(self, action: StepAction, **kwargs):
 
     @add_step_action.register(Callable)
     def _(self, action: StepActionFunc, **kwargs):
-        return self.add_step_action(StepAction(action), **kwargs)
+        return self.add_step_action(StepAction(action))
 
     @singledispatchmethod
-    def if_not_stopping(self, action: StepActionTypes, **kwargs) -> D:
+    def if_not_stopping(self, action: StepActionTypes) -> D:
         """
         Adds an action to be performed on each item that does not match any stop condition.
 
         :param action: The action to perform on each non-stopping item.
-        :keyword allow_re_wrapping: Allow rewrapping of :class:`StepAction`s with debug logging
         :return: The current traversal instance.
         """
         raise RuntimeError(f'StepAction [{action}] does not match expected: [StepAction | StepActionWithContextValue | Callable]')
 
     @if_not_stopping.register(Callable)
-    def _(self, action: StepActionFunc, **kwargs) -> D:
-        return self.add_step_action(lambda it, context: action(it, context) if not context.is_stopping else None, **kwargs)
+    def _(self, action: StepActionFunc) -> D:
+        return self.add_step_action(lambda it, context: action(it, context) if not context.is_stopping else None)
 
     @if_not_stopping.register
-    def _(self, action: StepAction, **kwargs) -> D:
+    def _(self, action: StepAction) -> D:
         action.apply = lambda it, context: action._func(it, context) if not context.is_stopping else None
-        return self.add_step_action(action, **kwargs)
+        return self.add_step_action(action)
 
     @singledispatchmethod
-    def if_stopping(self, action: StepActionTypes, **kwargs) -> D:
+    def if_stopping(self, action: StepActionTypes) -> D:
         """
         Adds an action to be performed on each item that matches a stop condition.
 
         :param action: The action to perform on each stopping item.
-        :keyword allow_re_wrapping: Allow rewrapping of :class:`StepActions`s with debug logging
         :return: The current traversal instance.
         """
         raise RuntimeError(f'StepAction [{action}] does not match expected: [StepAction | StepActionWithContextValue | Callable]')
 
     @if_stopping.register(Callable)
-    def _(self, action: StepActionFunc, **kwargs) -> D:
-        return self.add_step_action(lambda it, context: action(it, context) if context.is_stopping else None, **kwargs)
+    def _(self, action: StepActionFunc) -> D:
+        return self.add_step_action(lambda it, context: action(it, context) if context.is_stopping else None)
 
     @if_stopping.register
-    def _(self, action: StepAction, **kwargs) -> D:
+    def _(self, action: StepAction) -> D:
         action.apply = lambda it, context: action._func(it, context) if context.is_stopping else None
-        return self.add_step_action(action, **kwargs)
+        return self.add_step_action(action)
 
     def copy_step_actions(self, other: Traversal[T, D]) -> D:
         """

src/zepben/evolve/util.py

@@ -183,11 +183,6 @@ def __get__(self, cls, owner: T) -> T:
         return classmethod(self.fget).__get__(None, owner)()
 
 
-def extra_kwargs_not_allowed(kwargs, function_name):
-    if kwargs:
-        raise TypeError(f"'{kwargs.pop()}' is an invalid keyword argument for {function_name}()")
-
-
 def datetime_to_timestamp(date_time: datetime) -> PBTimestamp:
     timestamp = PBTimestamp()
     timestamp.FromDatetime(date_time)