DjangoGraphviz: modelviz.5.py

File modelviz.5.py, 6.3 KB (added by alexander@…, 7 years ago)

-i, --include_models only include selected models in graph

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