diff --git a/resources/graphs/state_machine/graph.py b/resources/graphs/state_machine/graph.py index 4e39cc5c0ed3e32b1b7c7d7c1628524e3a376383..e5e22b7ec4dc91979622737cbe0eaaaa59018c20 100644 --- a/resources/graphs/state_machine/graph.py +++ b/resources/graphs/state_machine/graph.py @@ -353,3 +353,22 @@ class StateMachineGraph(DirectedGraph): self._edges[connecting_edge.get_name()] = connecting_edge return connecting_edge + + def remove_node(self, node_indentifier): + """ + Removes node from graph. + :param node_indentifier: Node or name of node to remove. + :return: Removed node. + """ + node = self.get_node(node_indentifier) + + # Remove initial state, if present. + if node.is_initial: + self._initial_states.pop(node.get_name()) + + # Remove final state, if present. + if node.is_final: + self._final_states.pop(node.get_name()) + + # Call parent implementation. + return super().remove_node(node_indentifier)