DjangoGraphviz: modelviz.7.py

File modelviz.7.py, 6.7 KB (added by Christopher Schmidt <crschmidt@…>, 7 years ago)

Add exclude_models as well as include_models

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