path: root/umlpy.py

#!/usr/bin/env python
# coding=utf-8

'''
 umlpy
 Author : Guillaume "iXce" Seguin
 Email  : guillaume@segu.in (or guillaume.seguin@ens.fr)

 # UML-like graph produced for Python #

 Copyright (C) 2009 Guillaume Seguin

 This program is free software; you can redistribute it and/or
 modify it under the terms of the GNU General Public License
 as published by the Free Software Foundation; either version 2
 of the License, or (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 51 Franklin Street, Fifth Floor,
 Boston, MA  02110-1301, USA.
'''

import epydoc.docparser
import epydoc.docstringparser
import epydoc.apidoc

import gv

import sys
import os
import re

from optparse import OptionParser

parser = OptionParser ()
parser.add_option("-v", dest = "debug", action = "store_true",
                  default = False,
                  help = "enables debug output")
parser.add_option("-e", "--exclude", dest = "excludes", action = "append",
                  default = [],
                  help = "classes matching this regexp will be excluded \
                          from display")
parser.add_option("-i", "--include", dest = "includes", action = "append",
                  default = [],
                  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 \
                          display")
parser.add_option("--all-methods", dest = "all_methods",
                  action = "store_true", default = False,
                  help = "shows all methods")
parser.add_option("--all-properties", dest = "all_properties",
                  action = "store_true", default = False,
                  help = "shows all properties")
parser.add_option("--no-method", dest = "no_method",
                  action = "store_true", default = False,
                  help = "shows no method")
parser.add_option("--no-property", dest = "no_property",
                  action = "store_true", default = False,
                  help = "shows no property")
parser.add_option("-o", "--output", dest = "output", action = "store",
                  help = "output file")
parser.add_option("--png", dest = "format", action = "store_const",
                  const = "png", help = "output as png")
parser.add_option("--pdf", dest = "format", action = "store_const",
                  const = "pdf", default = "pdf", help = "output as pdf")
parser.add_option("--jpg", dest = "format", action = "store_const",
                  const = "jpg", help = "output as jpg")

DEFAULT_OUTPUT = "uml."

# Graph functions

CLASSNAME_FIELD_COLOR = "#FF6262"
PROPERTY_FIELD_COLOR = "#63BDFF"
METHOD_FIELD_COLOR = "#6EFF62"

SUPER_EDGE_COLOR = "#AD0006"
SUPER_EDGE_ARROWHEAD = "normal"

PROPERTY_EDGE_COLOR = "#002990"
PROPERTY_EDGE_ARROWHEAD = "diamond"

USE_EDGE_COLOR = "#05C800"
USE_EDGE_ARROWHEAD = "box"

LABEL_BASE = """<
<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0">
  <TR>
    <TD BGCOLOR="%s" PORT="%s">%s</TD>
  </TR>
%s
</TABLE>
>"""
FIELD_FORMATTER = """  <TR>
    <TD BGCOLOR="%s" PORT="%s">%s</TD>
  </TR>"""

class Graph (object):

    graph = None
    """@type: gv.digraph"""

    def __init__ (self):
        """@nodoc"""
        self.graph = gv.digraph ('g')
        gv.setv (self.graph, 'charset', 'utf-8')
        gv.setv (self.graph, 'overlap', 'false')
        gv.setv (self.graph, 'splines', 'true')
        gv.setv (self.graph, 'rankdir', 'BT')
        item = gv.protoedge (self.graph)
        gv.setv (item, 'len', '2')
        item = gv.protonode (self.graph)
        gv.setv (item, 'shape', 'plaintext')
        gv.setv (item, 'style', 'invis')

    def add_node (self, node):
        """@doc"""
        node.create (self.graph)

    def add_edge (self, edge):
        """@doc"""
        edge.create (self.graph)

    def render (self, filename, format):
        gv.layout (self.graph, 'dot')
        gv.render (self.graph, format, filename)

class Node (object):
    """@uses: umlpy.Graph"""

    class_name = None
    """@doc"""
    label = None
    """@doc"""

    def __init__ (self, class_name, properties, methods):
        """@nodoc"""
        self.class_name = class_name
        fields = [FIELD_FORMATTER % (PROPERTY_FIELD_COLOR, field, field)
                  for field in sorted (properties)]
        fields += [FIELD_FORMATTER % (METHOD_FIELD_COLOR, field, field)
                   for field in sorted (methods)]
        fields = "".join (fields)
        self.label = LABEL_BASE % (CLASSNAME_FIELD_COLOR,
                                   class_name,
                                   class_name,
                                   fields)
    
    def create (self, graph):
        """@doc"""
        node = gv.node (graph, self.class_name)
        gv.setv (node, 'label', self.label)
        return node

class SimpleNode (Node):

    def __init__ (self, class_name):
        """@nodoc"""
        super (SimpleNode, self).__init__ (class_name, [], [])
        self.label = LABEL_BASE % (CLASSNAME_FIELD_COLOR,
                                   class_name,
                                   class_name,
                                   "")

class Edge (object):
    """@uses: umlpy.Graph"""

    color = None
    """@doc"""
    arrowhead = None
    """@doc"""
    head = None
    """@doc"""
    tail = None
    """@doc"""

    def __init__ (self, head, tail):
        """@nodoc"""
        self.head = head
        self.tail = tail

    def create (self, graph):
        """@doc"""
        edge = gv.edge (graph, self.head, self.tail)
        gv.setv (edge, 'color', self.color)
        gv.setv (edge, 'arrowhead', self.arrowhead)
        return edge

class SuperEdge (Edge):

    color = SUPER_EDGE_COLOR
    arrowhead = SUPER_EDGE_ARROWHEAD

class UseEdge (Edge):

    color = USE_EDGE_COLOR
    arrowhead = USE_EDGE_ARROWHEAD

class PropertyEdge (Edge):

    color = PROPERTY_EDGE_COLOR
    arrowhead = PROPERTY_EDGE_ARROWHEAD

    property = None
    """@doc"""

    def __init__ (self, class_name, property_name, type_name):
        """@nodoc"""
        super (PropertyEdge, self).__init__ (type_name, class_name)
        self.property = property_name

    def create (self, graph):
        """@nodoc"""
        edge = super (PropertyEdge, self).create (graph)
        gv.setv (edge, 'headport', "%s:e" % self.property)
        return edge

class UmlPy (object):
    """@uses: umlpy.Graph"""

    targets = None
    """@doc"""

    excludes = None
    """@doc"""
    includes = None
    """@doc"""
    forces = None
    """@doc"""
    prefix = None
    """@doc"""
    all_methods = None
    """@doc"""
    all_properties = None
    """@doc"""
    no_method = None
    """@doc"""
    no_property = None
    """@doc"""
    output = None
    """@doc"""
    format = None
    """@doc"""

    debug = False
    """@doc"""

    classes = []
    bases_dict = {}
    uses_dict = {}
    methods_dict = {}
    vars_dict = {}

    def __init__ (self, targets, options):
        """@nodoc"""
        self.targets = targets
        self.excludes = reduce (lambda x, y: x + y,
                                map (lambda s: s.split (","), options.excludes),
                                [])
        self.includes = reduce (lambda x, y: x + y,
                                map (lambda s: s.split (","), options.includes),
                                [])
        self.forces = reduce (lambda x, y: x + y,
                              map (lambda s: s.split (","), options.forces),
                              [])
        self.real_excludes = [re.compile (exp)
                              for exp in self.excludes]
        self.real_includes = [re.compile (exp)
                              for exp in self.includes]
        self.real_forces = [re.compile (exp)
                            for exp in self.forces]
        self.prefix = options.prefix
        self.all_methods = options.all_methods
        self.all_properties = options.all_properties
        self.no_method = options.no_method
        self.no_property = options.no_property
        self.output = options.output
        self.format = options.format
        if not self.output:
            self.output = DEFAULT_OUTPUT + self.format
            pass
        self.debug = options.debug
        self.print_settings ()
        self.classes = []
        self.bases_dict = {}
        self.uses_dict = {}
        self.methods_dict = {}
        self.properties_dict = {}
        self.run ()

    def print_settings (self):
        """@nodoc"""
        print "Settings"
        print "--------"
        print "Targets :", ", ".join (self.targets)
        print "Output file :", self.output
        print "Prefix :", self.prefix
        print "Show all methods :", self.all_methods
        print "Show all properties :", self.all_properties
        print "Show no method :", self.no_method
        print "Show no propertie :", self.no_property
        print "Excludes :", self.excludes
        print "Includes :", self.includes
        print "Forces :", self.forces

    def is_excluded (self, class_name):
        """@nodoc"""
        for exp in self.real_excludes:
            if exp.match (class_name):
                for exp in [self.real_includes + self.real_forces]:
                    if exp.match (class_name):
                        return False
                return True
        else:
            return False

    def is_forced (self, class_name):
        """@nodoc"""
        for exp in self.real_forces:
            if exp.match (class_name):
                return True
        else:
            return False

    def run (self):
        """@nodoc"""
        docs = self.load_targets ()
        for doc in docs:
            self.process_doc (doc)
        self.render_graph ()

    def load_targets (self):
        """@nodoc"""
        paths = []
        for target in self.targets:
            for base in sys.path:
                path = os.path.join (base, target)
                path2 = os.path.join (path, "__init__.py")
                path3 = path + ".py"
                if os.path.isdir (path) and os.path.exists (path2):
                    paths.append ([os.path.join (path, file)
                                   for file in os.listdir (path)
                                   if file.endswith (".py")])
                    break
                elif os.path.exists (path3):
                    paths.append (path3)
                    break
        return [epydoc.docparser.parse_docs (path) for path in paths]

    def process_doc (self, doc):
        """@nodoc"""
        for var in doc.variables.values ():
            if type (var.value) != epydoc.apidoc.ClassDoc:
                continue
            self.process_class (var)

    def process_class (self, variable):
        """@nodoc"""
        value = variable.value
        name = str (value.canonical_name)
        name = name.replace (self.prefix, "", 1)
        if self.debug:
            print "Processing class", name
        self.classes.append (name)
        self.properties_dict[name] = []
        self.methods_dict[name] = []
        self.uses_dict[name] = []
        self.process_class_supers (name, value.bases)
        docstring = str (value.docstring)
        self.process_class_docstring (name, docstring)
        class_variables = value.variables.values ()
        for class_variable in class_variables:
            self.process_class_variable (name, class_variable)
    
    def process_class_supers (self, class_name, bases):
        """@nodoc"""
        bases = [str (base.canonical_name).replace (self.prefix, "", 1)
                 for base in bases
                 if str (base.canonical_name) not in ("object", "<UNKNOWN>")]
        if self.debug:
            if bases:
                print "Class %s has %s as supers" % (class_name,
                                                     ", ".join (bases))
            else:
                print "Class %s has no super" % class_name
        self.bases_dict[class_name] = bases

    def process_class_docstring (self, class_name, docstring):
        """@nodoc"""
        if docstring != "<UNKNOWN>":
            bits = docstring.split ("\n")
            for bit in bits:
                if bit.strip ().startswith ("@uses:"):
                    bit = str (bit.replace ("@uses:", "", 1).strip ())
                    self.uses_dict[class_name].append (bit)
                    if self.debug:
                        print "Class %s uses class %s" % (class_name, bit)

    def process_class_variable (self, class_name, variable):
        """@nodoc"""
        var_type = type (variable.value)
        if var_type == epydoc.apidoc.GenericValueDoc:
            if self.no_property:
                return
            self.process_class_property (class_name, variable)
        elif var_type in (epydoc.apidoc.RoutineDoc,
                          epydoc.apidoc.StaticMethodDoc):
            if self.no_method:
                return
            self.process_class_method (class_name, variable)
        else:
            raise ValueError, ("Unknown class variable type :" + var_type)

    @staticmethod
    def get_property_type (property):
        """@nodoc"""
        if property.type_descr != None:
            return str (property.type_descr.to_plaintext ("").strip ())
        else:
            return None

    def process_class_property (self, class_name, property):
        """@nodoc"""
        epydoc.docstringparser.parse_docstring (property, None)
        property_type = UmlPy.get_property_type (property)
        property_name = str (property.name.replace (self.prefix, "", 1))
        if self.debug:
            print "Found property %s of class %s of type %s" % \
                    (class_name, property_name, property_type)
        docstring = str (property.docstring)
        if (self.all_properties and not "@nodoc" in docstring) \
           or "@doc" in docstring or property_type:
            self.properties_dict[class_name].append ((property_name,
                                                      property_type))

    def process_class_method (self, class_name, method):
        """@nodoc"""
        docstring = str (method.value.docstring)
        method_name = str (method.name)
        if self.debug:
            print "Found method %s of class %s" % (class_name, method_name)
        if (self.all_methods and not "@nodoc" in docstring) \
           or "@doc" in docstring :
            self.methods_dict[class_name].append (method_name)

    def render_graph (self):
        other_classes = []
        graph = Graph ()
        for class_name in self.classes:
            if self.is_excluded (class_name):
                continue
            properties = self.properties_dict[class_name]
            property_names = [name for (name, type_name) in properties]
            property_types = [type_name
                              for (name, type_name) in properties
                              if type_name]
            methods = self.methods_dict[class_name]
            uses = self.uses_dict[class_name]
            bases = self.bases_dict[class_name]
            if properties or methods or uses or bases:
                graph.add_node (Node (class_name, property_names, methods))
            elif self.is_forced (class_name):
                graph.add_node (SimpleNode (class_name))
            for other_name in property_types + uses + bases:
                if not other_name or self.is_excluded (other_name):
                    continue
                if other_name not in other_classes \
                   and other_name not in self.classes:
                    other_classes.append (other_name)
                    graph.add_node (SimpleNode (other_name))
            for (property, type_name) in properties:
                if not type_name or self.is_excluded (type_name):
                    continue
                graph.add_edge (PropertyEdge (class_name, property, type_name))
            for base_name in bases:
                if self.is_excluded (base_name):
                    continue
                graph.add_edge (SuperEdge (class_name, base_name))
            for used_name in uses:
                if self.is_excluded (used_name):
                    continue
                graph.add_edge (UseEdge (class_name, used_name))
        graph.render (self.output, self.format)

if __name__ == "__main__":
    (options, args) = parser.parse_args ()
    UmlPy (args, options)