Code

DjangoGraphviz: modelviz.8.py

File modelviz.8.py, 7.5 KB (added by springmeyer, 5 years ago)

Added support for GeoDjango GeometryColumns and SpatialRefSys (requires fix in r9483)

Line 
1#!/usr/bin/env python
2"""Django model to DOT (Graphviz) converter
3by Antonio Cavedoni <antonio@cavedoni.org>
4
5Make sure your DJANGO_SETTINGS_MODULE is set to your project or
6place this script in the same directory of the project and call
7the script like this:
8
9$ python modelviz.py [-h] [-d] [-i <model_names>] [-e <model_names>] <app_label> ... <app_label> > <filename>.dot
10$ dot <filename>.dot -Tpng -o <filename>.png
11
12options:
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    -i, --include_models=User,Person,Car
20    only include selected models in graph.
21   
22    -e, --exclude_models=User,Person,Car
23    only include selected models in graph.
24"""
25__version__ = "0.8"
26__svnid__ = "$Id$"
27__license__ = "Python"
28__author__ = "Antonio Cavedoni <http://cavedoni.com/>"
29__contributors__ = [
30   "Stefano J. Attardi <http://attardi.org/>",
31   "limodou <http://www.donews.net/limodou/>",
32   "Carlo C8E Miron",
33   "Andre Campos <cahenan@gmail.com>",
34   "Justin Findlay <jfindlay@gmail.com>",
35   "Alexander Houben <alexander@houben.ch>",
36   "Christopher Schmidt <crschmidt@metacarta.com>",
37   "Dane Springmeyer <dane.springmeyer@gmail.com>"
38   ]
39
40import getopt, sys
41
42from django.core.management import setup_environ
43
44try:
45    import settings
46except ImportError:
47    pass
48else:
49    setup_environ(settings)
50
51from django.template import Template, Context
52from django.db import models
53from django.db.models import get_models
54from django.db.models.fields.related import \
55    ForeignKey, OneToOneField, ManyToManyField
56   
57from django.contrib.gis.db.models.fields import GeometryField
58
59try:
60    from django.db.models.fields.generic import GenericRelation
61except ImportError:
62    from django.contrib.contenttypes.generic import GenericRelation
63
64head_template = """
65digraph name {
66  fontname = "Helvetica"
67  fontsize = 8
68
69  node [
70    fontname = "Helvetica"
71    fontsize = 8
72    shape = "plaintext"
73  ]
74   edge [
75    fontname = "Helvetica"
76    fontsize = 8
77  ]
78
79"""
80
81body_template = """
82  {% for model in models %}
83    {% for relation in model.relations %}
84    {{ relation.target }} [label=<
85        <TABLE BGCOLOR="palegoldenrod" BORDER="0" CELLBORDER="0" CELLSPACING="0">
86        <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER" BGCOLOR="olivedrab4"
87        ><FONT FACE="Helvetica Bold" COLOR="white"
88        >{{ relation.target }}</FONT></TD></TR>
89        </TABLE>
90        >]
91    {{ model.name }} -> {{ relation.target }}
92    [label="{{ relation.name }}"] {{ relation.arrows }};
93    {% endfor %}
94  {% endfor %}
95
96  {% for model in models %}
97    {{ 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 }}</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 Bold">{{ field.name }}</FONT
107        ></TD>
108        <TD ALIGN="LEFT"
109        ><FONT {% if field.blank %}COLOR="#7B7B7B" {% endif %}FACE="Helvetica Bold">{{ field.type }}</FONT
110        ></TD></TR>
111        {% endfor %}
112    {% endif %}
113    </TABLE>
114    >]
115  {% endfor %}
116"""
117
118tail_template = """
119}
120"""
121
122def generate_dot(app_labels, **kwargs):
123    disable_fields = kwargs.get('disable_fields', False)
124    include_models = kwargs.get('include_models', [])
125    exclude_models = kwargs.get('exclude_models', [])
126
127    dot = head_template
128
129    for app_label in app_labels:
130        app = models.get_app(app_label)
131        graph = Context({
132            'name': '"%s"' % app.__name__,
133            'disable_fields': disable_fields,
134            'models': []
135            })
136
137        for appmodel in get_models(app):
138           
139            # consider given model name ?
140            def consider(model_name):
141                return (not include_models or model_name in include_models) and (not model_name in exclude_models)
142           
143            if not consider(appmodel._meta.object_name):
144                continue
145           
146            model = {
147                'name': appmodel.__name__,
148                'fields': [],
149                'relations': []
150                }
151
152            # model attributes
153            def add_attributes():
154                model['fields'].append({
155                    'name': field.name,
156                    'type': type(field).__name__,
157                    'blank': field.blank
158                    })
159
160            for field in appmodel._meta.fields:
161                add_attributes()
162
163            if appmodel._meta.many_to_many:
164                for field in appmodel._meta.many_to_many:
165                    add_attributes()
166
167            # relations
168            def add_relation(extras=""):
169                _rel = {
170                    'target': field.rel.to.__name__,
171                    'type': type(field).__name__,
172                    'name': field.name,
173                    'arrows': extras
174                    }
175                if _rel not in model['relations'] and consider(_rel['target']):
176                    model['relations'].append(_rel)
177
178            def add_geo_relation(target,relation,extras=""):
179                _rel = {
180                    'target': target,
181                    'type': type(field).__name__,
182                    'name': relation,
183                    'arrows': extras
184                    }
185                if _rel not in model['relations'] and consider(_rel['target']):
186                    model['relations'].append(_rel)
187
188            for field in appmodel._meta.fields:
189                if isinstance(field, ForeignKey):
190                    add_relation()
191                elif isinstance(field, OneToOneField):
192                    add_relation("[arrowhead=none arrowtail=none]")
193                elif isinstance(field, GeometryField):
194                    #import pdb; pdb.set_trace()
195                    add_geo_relation('SpatialRefSys',field._srid,"[arrowhead=normal arrowtail=none]")
196                    add_geo_relation('GeometryColumns',field._geom,"[arrowhead=normal arrowtail=none]")
197
198            if appmodel._meta.many_to_many:
199                for field in appmodel._meta.many_to_many:
200                    if isinstance(field, ManyToManyField):
201                        add_relation("[arrowhead=normal arrowtail=normal]")
202                    elif isinstance(field, GenericRelation):
203                        add_relation(
204                            '[style="dotted"] [arrowhead=normal arrowtail=normal]')
205            graph['models'].append(model)
206               
207        t = Template(body_template)
208        dot += '\n' + t.render(graph)
209
210    dot += '\n' + tail_template
211
212    return dot
213
214def main():
215    try:
216        opts, args = getopt.getopt(sys.argv[1:], "hdi:e:",
217                    ["help", "disable_fields", "include_models=", "exclude_models="])
218    except getopt.GetoptError, error:
219        print __doc__
220        sys.exit(error)
221    else:
222        if not args:
223            print __doc__
224            sys.exit()
225
226    kwargs = {}
227    for opt, arg in opts:
228        if opt in ("-h", "--help"):
229            print __doc__
230            sys.exit()
231        if opt in ("-d", "--disable_fields"):
232            kwargs['disable_fields'] = True
233        if opt in ("-i", "--include_models"):
234            kwargs['include_models'] = arg.split(',')
235        if opt in ("-e", "--exclude_models"):
236            kwargs['exclude_models'] = arg.split(',')
237    print generate_dot(args, **kwargs)
238
239if __name__ == "__main__":
240    main()