contrib/modelviz.py

   1 #!/usr/bin/env python
   2 """Django model to DOT (Graphviz) converter
   3 by Antonio Cavedoni <antonio@cavedoni.org>
   4
   5 Make sure your DJANGO_SETTINGS_MODULE is set to your project or
   6 place this script in the same directory of the project and call
   7 the script like this:
   8
   9 $ python modelviz.py [-h] [-d] <app_label> ... <app_label> > <filename>.dot
  10 $ dot <filename>.dot -Tpng -o <filename>.png
  11
  12 options:
  13     -h, --help
  14     show this help message and exit.
  15
  16     -d, --disable_fields
  17     don't show the class member fields.
  18 """
  19 __version__ = "0.8"
  20 __svnid__ = "$Id$"
  21 __license__ = "Python"
  22 __author__ = "Antonio Cavedoni <http://cavedoni.com/>"
  23 __contributors__ = [
  24    "Stefano J. Attardi <http://attardi.org/>",
  25    "limodou <http://www.donews.net/limodou/>",
  26    "Carlo C8E Miron",
  27    "Andre Campos <cahenan@gmail.com>",
  28    "Justin Findlay <jfindlay@gmail.com>",
  29    ]
  30
  31 import getopt, sys
  32
  33 from django.core.management import setup_environ
  34
  35 try:
  36     import settings
  37 except ImportError:
  38     pass
  39 else:
  40     setup_environ(settings)
  41
  42 from django.template import Template, Context
  43 from django.db import models
  44 from django.db.models import get_models
  45 from django.db.models.fields.related import \
  46     ForeignKey, OneToOneField, ManyToManyField
  47
  48 try:
  49     from django.db.models.fields.generic import GenericRelation
  50 except ImportError:
  51     from django.contrib.contenttypes.generic import GenericRelation
  52
  53 head_template = """
  54 digraph name {
  55   fontname = "Helvetica"
  56   fontsize = 8
  57
  58   node [
  59     fontname = "Helvetica"
  60     fontsize = 8
  61     shape = "plaintext"
  62   ]
  63    edge [
  64     fontname = "Helvetica"
  65     fontsize = 8
  66   ]
  67
  68 """
  69
  70 body_template = """
  71   {% for model in models %}
  72     {% for relation in model.relations %}
  73     {{ relation.target }} [label=<
  74         <TABLE BGCOLOR="palegoldenrod" BORDER="0" CELLBORDER="0" CELLSPACING="0">
  75         <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER" BGCOLOR="olivedrab4"
  76         ><FONT FACE="Helvetica Bold" COLOR="white"
  77         >{{ relation.target }}</FONT></TD></TR>
  78         </TABLE>
  79         >]
  80     {{ model.name }} -> {{ relation.target }}
  81     [label="{{ relation.name }}"] {{ relation.arrows }};
  82     {% endfor %}
  83   {% endfor %}
  84
  85   {% for model in models %}
  86     {{ model.name }} [label=<
  87     <TABLE BGCOLOR="palegoldenrod" BORDER="0" CELLBORDER="0" CELLSPACING="0">
  88      <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER" BGCOLOR="olivedrab4"
  89      ><FONT FACE="Helvetica Bold" COLOR="white"
  90      >{{ model.name }}</FONT></TD></TR>
  91
  92     {% if not disable_fields %}
  93         {% for field in model.fields %}
  94         <TR><TD ALIGN="LEFT" BORDER="0"
  95         ><FONT {% if field.blank %}COLOR="#7B7B7B" {% endif %}FACE="Helvetica Bold">{{ field.name }}</FONT
  96         ></TD>
  97         <TD ALIGN="LEFT"
  98         ><FONT {% if field.blank %}COLOR="#7B7B7B" {% endif %}FACE="Helvetica Bold">{{ field.type }}</FONT
  99         ></TD></TR>
 100         {% endfor %}
 101     {% endif %}
 102     </TABLE>
 103     >]
 104   {% endfor %}
 105 """
 106
 107 tail_template = """
 108 }
 109 """
 110
 111 def generate_dot(app_labels, **kwargs):
 112     disable_fields = kwargs.get('disable_fields', False)
 113
 114     dot = head_template
 115
 116     for app_label in app_labels:
 117         app = models.get_app(app_label)
 118         graph = Context({
 119             'name': '"%s"' % app.__name__,
 120             'disable_fields': disable_fields,
 121             'models': []
 122             })
 123
 124         for appmodel in get_models(app):
 125             model = {
 126                 'name': appmodel.__name__,
 127                 'fields': [],
 128                 'relations': []
 129                 }
 130
 131             # model attributes
 132             def add_attributes():
 133                 model['fields'].append({
 134                     'name': field.name,
 135                     'type': type(field).__name__,
 136                     'blank': field.blank
 137                     })
 138
 139             for field in appmodel._meta.fields:
 140                 add_attributes()
 141
 142             if appmodel._meta.many_to_many:
 143                 for field in appmodel._meta.many_to_many:
 144                     add_attributes()
 145
 146             # relations
 147             def add_relation(extras=""):
 148                 _rel = {
 149                     'target': field.rel.to.__name__,
 150                     'type': type(field).__name__,
 151                     'name': field.name,
 152                     'arrows': extras
 153                     }
 154                 if _rel not in model['relations']:
 155                     model['relations'].append(_rel)
 156
 157             for field in appmodel._meta.fields:
 158                 if isinstance(field, ForeignKey):
 159                     add_relation()
 160                 elif isinstance(field, OneToOneField):
 161                     add_relation("[arrowhead=none arrowtail=none]")
 162
 163             if appmodel._meta.many_to_many:
 164                 for field in appmodel._meta.many_to_many:
 165                     if isinstance(field, ManyToManyField):
 166                         add_relation("[arrowhead=normal arrowtail=normal]")
 167                     elif isinstance(field, GenericRelation):
 168                         add_relation(
 169                             '[style="dotted"] [arrowhead=normal arrowtail=normal]')
 170             graph['models'].append(model)
 171
 172         t = Template(body_template)
 173         dot += '\n' + t.render(graph)
 174
 175     dot += '\n' + tail_template
 176
 177     return dot
 178
 179 def main():
 180     try:
 181         opts, args = getopt.getopt(sys.argv[1:], "hd",
 182                     ["help", "disable_fields"])
 183     except getopt.GetoptError, error:
 184         print __doc__
 185         sys.exit(error)
 186     else:
 187         if not args:
 188             print __doc__
 189             sys.exit()
 190
 191     kwargs = {}
 192     for opt, arg in opts:
 193         if opt in ("-h", "--help"):
 194             print __doc__
 195             sys.exit()
 196         if opt in ("-d", "--disable_fields"):
 197             kwargs['disable_fields'] = True
 198     print generate_dot(args, **kwargs)
 199
 200 if __name__ == "__main__":
 201     main()