diff --git a/taskflow/types/fsm.py b/taskflow/types/fsm.py index df614767..afb4eb1d 100644 --- a/taskflow/types/fsm.py +++ b/taskflow/types/fsm.py @@ -23,6 +23,7 @@ import six from taskflow import exceptions as excp from taskflow.types import table +from taskflow.utils import misc class _Jump(object): @@ -97,6 +98,7 @@ class FSM(object): return False return self._states[self._current.name]['terminal'] + @misc.disallow_when_frozen(FrozenMachine) def add_state(self, state, terminal=False, on_enter=None, on_exit=None): """Adds a given state to the state machine. @@ -111,8 +113,6 @@ class FSM(object): :param state: state being entered or exited :type state: string """ - if self.frozen: - raise FrozenMachine() if state in self._states: raise excp.Duplicate("State '%s' already defined" % state) if on_enter is not None: @@ -129,6 +129,7 @@ class FSM(object): } self._transitions[state] = OrderedDict() + @misc.disallow_when_frozen(FrozenMachine) def add_reaction(self, state, event, reaction, *args, **kwargs): """Adds a reaction that may get triggered by the given event & state. @@ -149,8 +150,6 @@ class FSM(object): processed (and this process typically repeats) until the state machine reaches a terminal state. """ - if self.frozen: - raise FrozenMachine() if state not in self._states: raise excp.NotFound("Can not add a reaction to event '%s' for an" " undefined state '%s'" % (event, state)) @@ -162,6 +161,7 @@ class FSM(object): raise excp.Duplicate("State '%s' reaction to event '%s'" " already defined" % (state, event)) + @misc.disallow_when_frozen(FrozenMachine) def add_transition(self, start, end, event): """Adds an allowed transition from start -> end for the given event. @@ -169,8 +169,6 @@ class FSM(object): :param end: end of the transition :param event: event that caused the transition """ - if self.frozen: - raise FrozenMachine() if start not in self._states: raise excp.NotFound("Can not add a transition on event '%s' that" " starts in a undefined state '%s'" % (event, diff --git a/taskflow/types/tree.py b/taskflow/types/tree.py index c4273149..d5b4c12e 100644 --- a/taskflow/types/tree.py +++ b/taskflow/types/tree.py @@ -21,6 +21,8 @@ import os import six +from taskflow.utils import misc + class FrozenNode(Exception): """Exception raised when a frozen node is modified.""" @@ -98,9 +100,12 @@ class Node(object): n.freeze() self.frozen = True + @misc.disallow_when_frozen(FrozenNode) def add(self, child): - if self.frozen: - raise FrozenNode() + """Adds a child to this node (appends to left of existing children). + + NOTE(harlowja): this will also set the childs parent to be this node. + """ child.parent = self self._children.append(child) diff --git a/taskflow/utils/misc.py b/taskflow/utils/misc.py index 6c5c9de7..ec9eda80 100644 --- a/taskflow/utils/misc.py +++ b/taskflow/utils/misc.py @@ -227,6 +227,23 @@ def look_for(haystack, needles, extractor=None): return [needles[i] for (_hay_i, i) in sorted(matches)] +def disallow_when_frozen(excp_cls): + """Frozen checking/raising method decorator.""" + + def decorator(f): + + @six.wraps(f) + def wrapper(self, *args, **kwargs): + if self.frozen: + raise excp_cls() + else: + return f(self, *args, **kwargs) + + return wrapper + + return decorator + + def clamp(value, minimum, maximum, on_clamped=None): """Clamps a value to ensure its >= minimum and <= maximum.""" if minimum > maximum: