From 48ba47dfaff7da2a9ed636231ebd70e1056bb245 Mon Sep 17 00:00:00 2001
From: Brandon Rodriguez <brodriguez8774@gmail.com>
Date: Sun, 17 Nov 2019 05:49:17 -0500
Subject: [PATCH] Add handling for initial/final states in "remove_node" graph
function
---
resources/graphs/state_machine/graph.py | 19 +++++++++++++++++++
1 file changed, 19 insertions(+)
diff --git a/resources/graphs/state_machine/graph.py b/resources/graphs/state_machine/graph.py
index 4e39cc5..e5e22b7 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)
--
GitLab