web/lib/django_extensions/management/commands/graph_models.py
author ymh <ymh.work@gmail.com>
Wed, 20 Jan 2010 12:37:40 +0100
changeset 3 526ebd3988b0
permissions -rw-r--r--
replace pocketfilms occurence by blinkster
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
3
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     1
from django.core.management.base import BaseCommand, CommandError
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     2
from optparse import make_option
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     3
from django_extensions.management.modelviz import generate_dot
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     4
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     5
class Command(BaseCommand):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     6
    option_list = BaseCommand.option_list + (
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     7
        make_option('--disable-fields', '-d', action='store_true', dest='disable_fields',
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     8
            help='Do not show the class member fields'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
     9
        make_option('--group-models', '-g', action='store_true', dest='group_models',
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    10
            help='Group models together respective to there application'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    11
        make_option('--all-applications', '-a', action='store_true', dest='all_applications',
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    12
            help='Automaticly include all applications from INSTALLED_APPS'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    13
        make_option('--output', '-o', action='store', dest='outputfile',
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    14
            help='Render output file. Type of output dependend on file extensions. Use png or jpg to render graph to image.'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    15
        make_option('--layout', '-l', action='store', dest='layout', default='dot',
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    16
            help='Layout to be used by GraphViz for visualization. Layouts: circo dot fdp neato nop nop1 nop2 twopi'),
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    17
    )
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    18
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    19
    help = ("Creates a GraphViz dot file for the specified app names.  You can pass multiple app names and they will all be combined into a single model.  Output is usually directed to a dot file.")
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    20
    args = "[appname]"
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    21
    label = 'application name'
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    22
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    23
    requires_model_validation = True
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    24
    can_import_settings = True
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    25
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    26
    def handle(self, *args, **options):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    27
        if len(args) < 1 and not options['all_applications']:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    28
            raise CommandError("need one or more arguments for appname")
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    29
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    30
        dotdata = generate_dot(args, **options)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    31
        if options['outputfile']:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    32
            self.render_output(dotdata, **options)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    33
        else:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    34
            self.print_output(dotdata)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    35
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    36
    def print_output(self, dotdata):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    37
        print dotdata
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    38
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    39
    def render_output(self, dotdata, **kwargs):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    40
        try:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    41
            import pygraphviz
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    42
        except ImportError, e:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    43
            raise CommandError("need pygraphviz python module ( apt-get install python-pygraphviz )")
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    44
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    45
        vizdata = ' '.join(dotdata.split("\n")).strip()
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    46
        version = pygraphviz.__version__.rstrip("-svn")
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    47
        try:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    48
            if [int(v) for v in version.split('.')]<(0,36):
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    49
                # HACK around old/broken AGraph before version 0.36 (ubuntu ships with this old version)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    50
                import tempfile
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    51
                tmpfile = tempfile.NamedTemporaryFile()
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    52
                tmpfile.write(vizdata)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    53
                tmpfile.seek(0)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    54
                vizdata = tmpfile.name
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    55
        except ValueError:
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    56
            pass
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    57
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    58
        graph = pygraphviz.AGraph(vizdata)
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    59
        graph.layout(prog=kwargs['layout'])
526ebd3988b0 replace pocketfilms occurence by blinkster
ymh <ymh.work@gmail.com>
parents:
diff changeset
    60
        graph.draw(kwargs['outputfile'])