adding graph.remove_recovered() and graph.reset() methods

graphism/edge.py

@@ -68,3 +68,11 @@ def to_dict(self):
 
     def length(self):
         return self.__length(self)
+
+    def edge_to_tuple(self):
+        """
+        Returns a tuple containing the names of a given edge's parent and child nodes.
+
+        :rtype tuple: A two-item tuple containing the names of the edge's parent and child nodes.
+        """
+        return (self.parent.name(), self.child.name())

graphism/graph.py

@@ -255,28 +255,55 @@ def add_infected(self, node):
 
         """
         self.__infected[node.name()] = node
-
+
+    def add_susceptible(self, node):
+        """
+        Adds a node to the list of susceptible nodes
+
+        """
+        self.__susceptible[node.name()] = node
+
     def add_recovered(self, node):
         """
         Adds a node to the list of recovered nodes
 
         """
         self.__recovered[node.name()] = node
-        
+
     def remove_infected(self, node):
         """
         Removes a node from the list of infected nodes
 
         """
         del self.__infected[node.name()]
-            
+
     def remove_susceptible(self, node):
         """
         Removes a node from the list of susceptible nodes.
 
         """
         del self.__susceptible[node.name()]
 
+    def remove_recovered(self, node):
+        """
+        Removes a node from the list of recovered nodes
+
+        """
+        del self.__recovered[node.name()]
+
+    def reset(self):
+        """
+        Resets graph to a state where all nodes are susceptible.
+
+        """
+        for node in self.recovered():
+            self.add_susceptible(node)
+            self.remove_recovered(node)
+
+        for node in self.infected():
+            self.add_susceptible(node)
+            self.remove_infected(node)
+
     def nodes(self):
         """
         Returns the internal nodes as a set. Deprecated.