Code

DjangoGraphviz: modelviz.3.py

File modelviz.3.py, 5.4 KB (added by André Campos, 7 years ago)

added other features.

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