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] [-a] [-d] [-g] [-i <model_names>] <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     -a, --all_applications

    17     show models from all applications.

    18 

    19     -d, --disable_fields

    20     don't show the class member fields.

    21 

    22     -g, --group_models

    23     draw an enclosing box around models from the same app.

    24 

    25     -i, --include_models=User,Person,Car

    26     only include selected models in graph.

    27 """

    28 __version__ = "0.9"

    29 __svnid__ = "$Id$"

    30 __license__ = "Python"

    31 __author__ = "Antonio Cavedoni <http://cavedoni.com/>"

    32 __contributors__ = [

    33    "Stefano J. Attardi <http://attardi.org/>",

    34    "limodou <http://www.donews.net/limodou/>",

    35    "Carlo C8E Miron",

    36    "Andre Campos <cahenan@gmail.com>",

    37    "Justin Findlay <jfindlay@gmail.com>",

    38    "Alexander Houben <alexander@houben.ch>",

    39    "Bas van Oostveen <v.oostveen@gmail.com>",

    40 ]

    41 

    42 import getopt, sys

    43 

    44 from django.core.management import setup_environ

    45 

    46 try:

    47     import settings

    48 except ImportError:

    49     pass

    50 else:

    51     setup_environ(settings)

    52 

    53 from django.utils.safestring import mark_safe

    54 from django.template import Template, Context

    55 from django.db import models

    56 from django.db.models import get_models

    57 from django.db.models.fields.related import \

    58     ForeignKey, OneToOneField, ManyToManyField

    59 

    60 try:

    61     from django.db.models.fields.generic import GenericRelation

    62 except ImportError:

    63     from django.contrib.contenttypes.generic import GenericRelation

    64 

    65 head_template = """

    66 digraph name {

    67   fontname = "Helvetica"

    68   fontsize = 8

    69 

    70   node [

    71     fontname = "Helvetica"

    72     fontsize = 8

    73     shape = "plaintext"

    74   ]

    75   edge [

    76     fontname = "Helvetica"

    77     fontsize = 8

    78   ]

    79 

    80 """

    81 

    82 body_template = """

    83 {% if use_subgraph %}

    84 subgraph {{ cluster_app_name }} {

    85   label=<

    86         <TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0">

    87         <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER"

    88         ><FONT FACE="Helvetica Bold" COLOR="Black" POINT-SIZE="12"

    89         >{{ app_name }}</FONT></TD></TR>

    90         </TABLE>

    91         >

    92   color=olivedrab4

    93   style="rounded"

    94 {% endif %}

    95 

    96   {% for model in models %}

    97     {{ model.app_name }}_{{ model.name }} [label=<

    98     <TABLE BGCOLOR="palegoldenrod" BORDER="0" CELLBORDER="0" CELLSPACING="0">

    99      <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER" BGCOLOR="olivedrab4"

   100      ><FONT FACE="Helvetica Bold" COLOR="white"

   101      >{{ model.name }}{% if model.abstracts %}<BR/>&lt;<FONT FACE="Helvetica Italic">{{ model.abstracts|join:"," }}</FONT>&gt;{% endif %}</FONT></TD></TR>

   102 

   103     {% if not disable_fields %}

   104         {% for field in model.fields %}

   105         <TR><TD ALIGN="LEFT" BORDER="0"

   106         ><FONT {% if field.blank %}COLOR="#7B7B7B" {% endif %}FACE="Helvetica {% if field.abstract %}Italic{% else %}Bold{% endif %}">{{ field.name }}</FONT

   107         ></TD>

   108         <TD ALIGN="LEFT"

   109         ><FONT {% if field.blank %}COLOR="#7B7B7B" {% endif %}FACE="Helvetica {% if field.abstract %}Italic{% else %}Bold{% endif %}">{{ field.type }}</FONT

   110         ></TD></TR>

   111         {% endfor %}

   112     {% endif %}

   113     </TABLE>

   114     >]

   115   {% endfor %}

   116 

   117 {% if use_subgraph %}

   118 }

   119 {% endif %}

   120 """

   121 

   122 rel_template = """

   123   {% for model in models %}

   124     {% for relation in model.relations %}

   125     {% if relation.needs_node %}

   126     {{ relation.target_app }}_{{ relation.target }} [label=<

   127         <TABLE BGCOLOR="palegoldenrod" BORDER="0" CELLBORDER="0" CELLSPACING="0">

   128         <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER" BGCOLOR="olivedrab4"

   129         ><FONT FACE="Helvetica Bold" COLOR="white"

   130         >{{ relation.target }}</FONT></TD></TR>

   131         </TABLE>

   132         >]

   133     {% endif %}

   134     {{ model.app_name }}_{{ model.name }} -> {{ relation.target_app }}_{{ relation.target }}

   135     [label="{{ relation.name }}"] {{ relation.arrows }};

   136     {% endfor %}

   137   {% endfor %}

   138 """

   139 

   140 tail_template = """

   141 }

   142 """

   143 

   144 def generate_dot(app_labels, **kwargs):

   145     disable_fields = kwargs.get('disable_fields', False)

   146     include_models = kwargs.get('include_models', [])

   147     all_applications = kwargs.get('all_applications', False)

   148     use_subgraph = kwargs.get('group_models', False)

   149 

   150     dot = head_template

   151 

   152     apps = []

   153     if all_applications:

   154         apps = models.get_apps()

   155 

   156     for app_label in app_labels:

   157         app = models.get_app(app_label)

   158         if not app in apps:

   159             apps.append(app)

   160 

   161     graphs = []

   162     for app in apps:

   163         graph = Context({

   164             'name': '"%s"' % app.__name__,

   165             'app_name': "%s" % '.'.join(app.__name__.split('.')[:-1]),

   166             'cluster_app_name': "cluster_%s" % app.__name__.replace(".", "_"),

   167             'disable_fields': disable_fields,

   168             'use_subgraph': use_subgraph,

   169             'models': []

   170         })

   171 

   172         for appmodel in get_models(app):

   173             abstracts = [e.__name__ for e in appmodel.__bases__ if hasattr(e, '_meta') and e._meta.abstract]

   174             abstract_fields = []

   175             for e in appmodel.__bases__:

   176                 if hasattr(e, '_meta') and e._meta.abstract:

   177                     abstract_fields.extend(e._meta.fields)

   178             model = {

   179                 'app_name': app.__name__.replace(".", "_"),

   180                 'name': appmodel.__name__,

   181                 'abstracts': abstracts,

   182                 'fields': [],

   183                 'relations': []

   184             }

   185 

   186             # consider given model name ?

   187             def consider(model_name):

   188                 return not include_models or model_name in include_models

   189 

   190             if not consider(appmodel._meta.object_name):

   191                 continue

   192 

   193             # model attributes

   194             def add_attributes(field):

   195                 model['fields'].append({

   196                     'name': field.name,

   197                     'type': type(field).__name__,

   198                     'blank': field.blank,

   199                     'abstract': field in abstract_fields,

   200                 })

   201 

   202             for field in appmodel._meta.fields:

   203                 add_attributes(field)

   204 

   205             if appmodel._meta.many_to_many:

   206                 for field in appmodel._meta.many_to_many:

   207                     add_attributes(field)

   208 

   209             # relations

   210             def add_relation(field, extras=""):

   211                 _rel = {

   212                     'target_app': field.rel.to.__module__.replace('.','_'),

   213                     'target': field.rel.to.__name__,

   214                     'type': type(field).__name__,

   215                     'name': field.name,

   216                     'arrows': extras,

   217                     'needs_node': True

   218                 }

   219                 if _rel not in model['relations'] and consider(_rel['target']):

   220                     model['relations'].append(_rel)

   221 

   222             for field in appmodel._meta.fields:

   223                 if isinstance(field, ForeignKey):

   224                     add_relation(field)

   225                 elif isinstance(field, OneToOneField):

   226                     add_relation(field, '[arrowhead=none arrowtail=none]')

   227 

   228             if appmodel._meta.many_to_many:

   229                 for field in appmodel._meta.many_to_many:

   230                     if isinstance(field, ManyToManyField) and getattr(field, 'creates_table', False):

   231                         add_relation(field, '[arrowhead=normal arrowtail=normal]')

   232                     elif isinstance(field, GenericRelation):

   233                         add_relation(field, mark_safe('[style="dotted"] [arrowhead=normal arrowtail=normal]'))

   234             graph['models'].append(model)

   235         graphs.append(graph)

   236 

   237     nodes = []

   238     for graph in graphs:

   239         nodes.extend([e['name'] for e in graph['models']])

   240 

   241     for graph in graphs:

   242         # don't draw duplication nodes because of relations

   243         for model in graph['models']:

   244             for relation in model['relations']:

   245                 if relation['target'] in nodes:

   246                     relation['needs_node'] = False

   247         # render templates

   248         t = Template(body_template)

   249         dot += '\n' + t.render(graph)

   250 

   251     for graph in graphs:

   252         t = Template(rel_template)

   253         dot += '\n' + t.render(graph)

   254 

   255     dot += '\n' + tail_template

   256     return dot

   257 

   258 def main():

   259     try:

   260         opts, args = getopt.getopt(sys.argv[1:], "hadgi:",

   261                     ["help", "all_applications", "disable_fields", "group_models", "include_models="])

   262     except getopt.GetoptError, error:

   263         print __doc__

   264         sys.exit(error)

   265     

   266     kwargs = {}

   267     for opt, arg in opts:

   268         if opt in ("-h", "--help"):

   269             print __doc__

   270             sys.exit()

   271         if opt in ("-a", "--all_applications"):

   272             kwargs['all_applications'] = True

   273         if opt in ("-d", "--disable_fields"):

   274             kwargs['disable_fields'] = True

   275         if opt in ("-g", "--group_models"):

   276             kwargs['group_models'] = True

   277         if opt in ("-i", "--include_models"):

   278             kwargs['include_models'] = arg.split(',')

   279 

   280     if not args and not kwargs.get('all_applications', False):

   281         print __doc__

   282         sys.exit()

   283 

   284     print generate_dot(args, **kwargs)

   285 

   286 if __name__ == "__main__":

   287     main()