small bugfix and new argument : subcat - use subcategories instead of categories

src/graph_builder.py

@@ -9,14 +9,19 @@ from operator import itemgetter
 import graph_helpers as gh
 import logging
 
-subcat=False
 
 
-def build_graph(years=[],include_utterances=False):
+def build_graph(years=[],include_utterances=False,subcat=False):
     
+
+    def maybe_subcat(string):
+        if not subcat:
+            return False
+        else: return string
+
     # we intialize an empty graph
     G = nx.MultiGraph()
-    
+     
     #index to keep track of nodes
     nodeindex=0
 
@@ -51,7 +56,7 @@ def build_graph(years=[],include_utterances=False):
         descriptions = [dh.HasDescriptionNode(of) for of in offs]
         logging.info("descriptions found: {}".format(descriptions))
         #[desc.simplify_to_direct_object() for desc in descriptions]
-        offs = [dh.OffenceNode(of,"offence",subcat=subcat,dummy=False) for of in offs]
+        offs = [dh.OffenceNode(of,"offence",subcat=maybe_subcat("offence"),dummy=False) for of in offs]
         #print(descriptions)
         #asd
         for i,off in enumerate(offs):
@@ -69,10 +74,10 @@ def build_graph(years=[],include_utterances=False):
         
         #similar to offences
         vs = div.findAll("rs",{"type":"verdictDescription"})
-        vs = [dh.VerdictNode(v,"verdict",subcat=subcat,dummy=False) for v in vs]
+        vs = [dh.VerdictNode(v,"verdict",subcat=maybe_subcat("verdict"),dummy=False) for v in vs]
         if not vs:
             logging.warning("warning no verdict found inserting dummy")
-            vn=dh.VerdictNode({},"",subcat=subcat,dummy=True,text=False)
+            vn=dh.VerdictNode({},"",subcat=maybe_subcat("dummy"),dummy=True,text=False)
             nodeindex,vindex = gh.maybe_new_index(vn,node_index_dict)
             G.add_node(vindex, label="Verdict", nodeobj=vn)
             G.add_edge(trialindex,vindex,edge_class="with-verdict")
@@ -85,7 +90,7 @@ def build_graph(years=[],include_utterances=False):
         
         # similar to offences
         puns = div.findAll("rs",{"type":"punishmentDescription"})
-        puns = [dh.PunishmentNode(pun,"punishment",subcat=subcat,dummy=False) for pun in puns]
+        puns = [dh.PunishmentNode(pun,"punishment",subcat=maybe_subcat("punishment"),dummy=False) for pun in puns]
         for pun in puns:
             
             nodeindex,punindex = gh.maybe_new_index(pun,node_index_dict)