Optionify prefix, add use relationship, show methods, rework code

1 files changed, 103 insertions, 53 deletions

diff --git a/grapher.py b/grapher.py
index 6e7629d..a9e7f19 100644
--- a/grapher.py
+++ b/grapher.py
@@ -23,6 +23,9 @@ parser.add_option("-i", "--include", dest = "includes", action = "append",
                   help = "classes matching this regexp will be included \
                           in display even if they were excluded by exclude \
                           regexps")
+parser.add_option("-p", "--prefix", dest = "prefix", action = "store",
+                  default = "",
+                  help = "prefix which will be stripped of class names")
 parser.add_option("-f", "--force", dest = "forces", action = "append",
                   default = [],
                   help = "classes matching this regexp will be forced into \
@@ -36,7 +39,11 @@ includes = reduce (lambda x, y: x + y,
                    map (lambda s: s.split(","), options.includes), [])
 forces = reduce (lambda x, y: x + y,
                  map (lambda s: s.split(","), options.forces), [])
+prefix = options.prefix
 
+print "Settings"
+print "--------"
+print "Prefix :", prefix
 print "Excludes :", excludes
 print "Includes :", includes
 print "Forces :", forces
@@ -80,6 +87,8 @@ for path in args:
 
 classes = []
 bases_dict = {}
+uses_dict = {}
+methods_dict = {}
 vars_dict = {}
 
 for doc in docs:
@@ -88,21 +97,34 @@ for doc in docs:
       continue
     var_val = var.value
     var_name = str (var_val.canonical_name)
-    var_name = var_name.replace ("netnavalbattle.", "")
+    var_name = var_name.replace (prefix, "")
     classes.append (var_name)
     if options.debug:
-        print var_name
+      print var_name
     var_vars = var_val.variables.values ()
     vars_dict[var_name] = []
+    methods_dict[var_name] = []
+    uses_dict[var_name] = []
+    if str (var_val.docstring) != "<UNKNOWN>":
+      bits = str (var_val.docstring).split ("\n")
+      for bit in bits:
+        if bit.startswith ("@uses:"):
+          bit = bit.replace ("@uses:", "").strip ()
+          uses_dict[var_name].append (bit)
     for var_var in var_vars:
-      if type (var_var.value) != epydoc.apidoc.GenericValueDoc:
-        continue
-      epydoc.docstringparser.parse_docstring (var_var, None)
-      if options.debug:
-        print var_name, var_var.name, get_var_type (var_var)
-      var_var_name = str (var_var.name.replace ("netnavalbattle.", ""))
-      vars_dict[var_name].append ((var_var_name, get_var_type (var_var)))
-    bases = [str (base.canonical_name).replace ("netnavalbattle.", "")
+      if type (var_var.value) == epydoc.apidoc.GenericValueDoc:
+        epydoc.docstringparser.parse_docstring (var_var, None)
+        if options.debug:
+          print var_name, var_var.name, get_var_type (var_var)
+        var_var_name = str (var_var.name.replace (prefix, ""))
+        vars_dict[var_name].append ((var_var_name, get_var_type (var_var)))
+      elif type (var_var.value) in (epydoc.apidoc.RoutineDoc,
+                                    epydoc.apidoc.StaticMethodDoc):
+        if "@doc" in str (var_var.value.docstring):
+          methods_dict[var_name].append (str (var_var.name))
+      else:
+        print type (var_var.value)
+    bases = [str (base.canonical_name).replace (prefix, "")
              for base in var_val.bases
              if str (base.canonical_name) not in ("object", "<UNKNOWN>")]
     if options.debug:
@@ -110,71 +132,83 @@ for doc in docs:
     bases_dict[var_name] = bases
 
 nodes_dict = {}
-fields_dict = {}
+var_fields_dict = {}
+method_fields_dict = {}
 graph = gv.digraph ('g')
 gv.setv (graph, 'charset', 'utf-8')
 gv.setv (graph, 'overlap', 'false')
 gv.setv (graph, 'splines', 'true')
+gv.setv (graph, 'rankdir', 'BT')
 item = gv.protoedge (graph)
 gv.setv (item, 'len', '2')
 item = gv.protonode (graph)
-gv.setv (item, 'shape', 'box')
-gv.setv (item, 'style', 'filled')
-
-def get_class_node (name):
-  if name not in nodes_dict:
-    node = gv.node (graph, name)
-    gv.setv (node, 'fillcolor', '#FF6262')
-    gv.setv (node, 'label', name)
-    fields_dict[name] = []
-    nodes_dict[name] = node
-  return nodes_dict[name]
+gv.setv (item, 'shape', 'plaintext')
 
 CLASS_COLOR = "#FF6262"
-FIELD_COLOR = "#63BDFF"
+VAR_FIELD_COLOR = "#63BDFF"
+METHOD_FIELD_COLOR = "#6EFF62"
 
-def build_record_label (class_name):
-    label_base = """<
+SUPER_EDGE_COLOR = "#AD0006"
+VAR_EDGE_COLOR = "#002990"
+USE_EDGE_COLOR = "#05C800"
+
+LABEL_BASE = """<
 <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0">
-    <TR>
-        <TD BGCOLOR="%s" PORT="%s">%s</TD>
-    </TR>
+  <TR>
+    <TD BGCOLOR="%s" PORT="%s">%s</TD>
+  </TR>
 %s
 </TABLE>
 >"""
-    field_formatter = """       <TR>
-        <TD BGCOLOR="%s" PORT="%s">%s</TD>
-    </TR>"""
-    fields = "".join ([field_formatter % (FIELD_COLOR, field, field)
-                      for field in fields_dict[class_name]])
-    return label_base % (CLASS_COLOR, class_name, class_name, fields)
-
-def add_var_field (class_name, var_name):
-  class_node = get_class_node (class_name)
-  if var_name not in fields_dict[class_name]:
-    fields_dict[class_name].append (var_name)
-    if gv.getv (class_node, 'shape') != 'plaintext':
-      gv.setv (class_node, 'shape', 'plaintext')
-      gv.setv (class_node, 'style', 'invis')
-    gv.setv (class_node, 'label', build_record_label (class_name))
+FIELD_FORMATTER = """  <TR>
+    <TD BGCOLOR="%s" PORT="%s">%s</TD>
+  </TR>"""
+
+def build_record_label (class_name):
+  fields = "".join ([FIELD_FORMATTER % (VAR_FIELD_COLOR, field, field)
+                    for field in var_fields_dict[class_name]])
+  fields += "".join ([FIELD_FORMATTER % (METHOD_FIELD_COLOR, field, field)
+                     for field in method_fields_dict[class_name]])
+  return LABEL_BASE % (CLASS_COLOR, class_name, class_name, fields)
+
+def build_simple_label (class_name):
+  return LABEL_BASE % (CLASS_COLOR, class_name, class_name, "")
+
+def check_class_node (name):
+  if name not in nodes_dict:
+    node = gv.node (graph, name)
+    gv.setv (node, 'shape', 'plaintext')
+    gv.setv (node, 'style', 'invis')
+    if name in var_fields_dict or name in method_fields_dict:
+      gv.setv (node, 'label', build_record_label (name))
+    else:
+      gv.setv (node, 'label', build_simple_label (name))
+    nodes_dict[name] = node
+  return nodes_dict[name]
 
 def add_super_edge (class_name, base_name):
-  class_node = get_class_node (class_name)
-  base_node = get_class_node (base_name)
+  check_class_node (class_name)
+  check_class_node (base_name)
   edge = gv.edge (graph, class_name, base_name)
   gv.setv (edge, 'arrowhead', 'normal')
-  if gv.getv (base_node, 'shape') == 'record':
-    gv.setv (edge, 'headport', "%s" % base_name)
+  gv.setv (edge, 'color', SUPER_EDGE_COLOR)
   # FIXME: edge style
 
 def add_var_edge (class_name, var_name, type_name):
-  if type_name == None:
-    return
-  add_var_field (class_name, var_name)
-  type_node = get_class_node (type_name)
+  check_class_node (class_name)
+  check_class_node (type_name)
   edge = gv.edge (graph, type_name, class_name)
   gv.setv (edge, 'headport', "%s:e" % var_name)
   gv.setv (edge, 'arrowhead', 'diamond')
+  gv.setv (edge, 'color', VAR_EDGE_COLOR)
+  # FIXME: edge style
+
+def add_use_edge (class_name, base_name):
+  check_class_node (class_name)
+  check_class_node (base_name)
+  edge = gv.edge (graph, class_name, base_name)
+  gv.setv (edge, 'arrowhead', 'box')
+  gv.setv (edge, 'color', USE_EDGE_COLOR)
   # FIXME: edge style
 
 for class_name in classes:
@@ -183,20 +217,36 @@ for class_name in classes:
   for (var_name, type_name) in vars_dict[class_name]:
     if not type_name or is_excluded (type_name):
       continue
-    add_var_edge (class_name, var_name, type_name)
+    if class_name not in var_fields_dict:
+        var_fields_dict[class_name] = []
+        method_fields_dict[class_name] = []
+    var_fields_dict[class_name].append (var_name)
+  for method_name in methods_dict[class_name]:
+    if class_name not in var_fields_dict:
+        var_fields_dict[class_name] = []
+        method_fields_dict[class_name] = []
+    method_fields_dict[class_name].append (method_name)
 
 for class_name in classes:
   if is_excluded (class_name):
     continue
+  for (var_name, type_name) in vars_dict[class_name]:
+    if not type_name or is_excluded (type_name):
+      continue
+    add_var_edge (class_name, var_name, type_name)
   for base_name in bases_dict[class_name]:
     if is_excluded (base_name):
       continue
     add_super_edge (class_name, base_name)
+  for used_name in uses_dict[class_name]:
+    if is_excluded (used_name):
+      continue
+    add_use_edge (class_name, used_name)
 
 for class_name in classes:
     for exp in forces:
         if exp.match (class_name):
-            get_class_node (class_name)
+            check_class_node (class_name)
             break
 
 gv.layout (graph, 'dot')