NF: method in GraphPArtition to check that index is a dict containing sets

Author: CindeeM

brainx/modularity.py

@@ -72,8 +72,12 @@ def __init__(self, graph, index):
         """
         # Store references to the original graph and label dict
         self.index = copy.deepcopy(index)
-        #self.graph = graph
 
+        ## add quick check to make sure the passed index is
+        ## a dict of sets
+        self._check_index_contains_sets()
+
+        
         # We'll need the graph's adjacency matrix often, so store it once
         self.graph_adj_matrix = nx.adj_matrix(graph)
 
@@ -102,6 +106,15 @@ def copy(self):
     def __len__(self):
         return len(self.index)
 
+
+    def _check_index_contains_sets(self):
+        """ the index in a GraphPartition is a dict of node sets
+        validate that the values of this dict are all of type(set)"""
+        index_types = [ type(x) for x in self.index.values() ]
+        if not all([ x== type(set()) for x in index_types]):
+            raise TypeError('index values should be of type set():: %s'%(index_types))
+
+
     def _edge_info(self, mod_e=None, mod_a=None, index=None):
         """Create the vectors of edge information.
 

brainx/tests/test_modularity.py

@@ -752,7 +752,7 @@ def test_adjust_partition():
 def test_empty_graphpartition():
     g = nx.Graph()
     g.add_node(1)
-    npt.assert_raises(ValueError, mod.GraphPartition, g, {1: g.nodes()})
+    npt.assert_raises(ValueError, mod.GraphPartition, g, {1: set(g.nodes())})
 
 
 if __name__ == "__main__":