Converting Django models into Graphviz DOT files

Inspired by ​this nice hack by Matt Biddulph I set out to create a similar thing for Django models. Below is an initial implementation, and here are some results generated with the OS X version of Graphviz:

• Graphviz DOT files for the Zyons project: ​comment ​event ​forum ​openid ​prefs ​sessions ​tag
• Resulting graphs in PDF: ​comment ​event ​forum ​openid ​prefs ​sessions ​tag

• Other samples:
  • ​generated dot file: resulting image in ​PDF, PNG
  • ​generated dot file: resulting image in ​PDF, PNG

You might also be interested in ​this Django app by Andrew Barilla from which I borrowed some ideas, that displays the graphviz results directly from the web.

#!/usr/bin/python
"""Django model to DOT (Graphviz) converter
by Antonio Cavedoni <antonio@cavedoni.org>

Make sure your DJANGO_SETTINGS_MODULE is set to your project and call 
the script like this:

  $ python modelviz.py <app_label> > <filename>.dot

Changelog

0.5.1
Got rid of print statements, now the generate_dot() function returns
a string with the file contents

0.5
Cleaned up code, now the relationship arrows start from the 
correct model attribute position

0.4
Fixed OneToOneField support (thanks, limodou)

0.3
Added support for GenericRelation and OneToOneField

0.2
Changed display (shape="record", thanks Malcolm Tredinnick), 
fixed arrow direction, added display of model attributes 
(thanks, Russell Keith-Magee)

0.1
First release.
"""

__version__ = "0.5.1"
__svnid__ = "$Id: modelviz.py 4 2006-08-06 19:48:42Z verbosus $"
__license__ = "Python"
__author__ = "Antonio Cavedoni <http://cavedoni.com/>"
__contributors__ = [
   "Stefano J. Attardi <http://attardi.org/>",
   "limodou <http://www.donews.net/limodou/>"
   ]

from django.db import models
from django.db.models import get_models
from django.db.models.fields.related import \
    ForeignKey, OneToOneField, ManyToManyField
from django.db.models.fields.generic import GenericRelation

def generate_dot(app_label):
   app = models.get_app(app_label)

   graph = []
   graph.append("digraph %s {" % ('"' + app.__name__ + '"'))
   graph.append("""  fontname = "Helvetica"
  fontsize = 8
      
  node [
    fontname = "Helvetica"
    fontsize = 8
    shape = "record"
  ]
   edge [
    fontname = "Helvetica"
    fontsize = 8
  ]
""")
   for o in get_models(app):
      graph.append(""" subgraph cluster_%(model)s {
  shape = "record";
  label = "%(model)s";
  fontname = "Helvetica Bold";
  fontsize = 10;
  labelfontcolor = "black";
  %(model)s [label = "{""" % {'model': o.__name__})

      # model attributes
      def add_attributes():
         graph.append("<%(model)s_%(field)s>%(field)s : %(related)s|" % \
                      {'model': o.__name__, 'field': field.name, 
                       'related': type(field).__name__})

      for field in o._meta.fields:
         add_attributes()
         
      if o._meta.many_to_many:
         for field in o._meta.many_to_many:
            add_attributes()

      graph.append(""" }"] [color="white" shape="record"];\n }\n""")

      # relations
      rel = []
      def add_relation(extras=""):
         _rel = """ %(model)s:%(model)s_%(field)s -> %(related)s 
         [label="%(relationship)s"] %(extras)s;\n""" % {
            'model': o.__name__, 'field': field.name, 
            'related': field.rel.to.__name__, 
            'relationship': type(field).__name__,
            'extras': extras}
         if _rel not in rel:
            rel.append(_rel)

      for field in o._meta.fields:
         if isinstance(field, ForeignKey):
            add_relation()
         elif isinstance(field, OneToOneField):
            add_relation("[arrowhead=none arrowtail=none]")

      if o._meta.many_to_many:
         for field in o._meta.many_to_many:
            if isinstance(field, ManyToManyField):
               add_relation("[arrowhead=normal arrowtail=normal]")
            elif isinstance(field, GenericRelation):
               add_relation(
                  '[style="dotted"] [arrowhead=normal arrowtail=normal]')
      graph.append("".join(rel))

   graph.append('}')
   return "".join(graph)

if __name__ == "__main__":
   import sys
   app_label = sys.argv[1]
   print generate_dot(app_label)