1 from django.core.management.base import BaseCommand, CommandError

     2 from optparse import make_option

     3 from django_extensions.management.modelviz import generate_dot

     4 

     5 class Command(BaseCommand):

     6     option_list = BaseCommand.option_list + (

     7         make_option('--disable-fields', '-d', action='store_true', dest='disable_fields',

     8             help='Do not show the class member fields'),

     9         make_option('--group-models', '-g', action='store_true', dest='group_models',

    10             help='Group models together respective to there application'),

    11         make_option('--all-applications', '-a', action='store_true', dest='all_applications',

    12             help='Automaticly include all applications from INSTALLED_APPS'),

    13         make_option('--output', '-o', action='store', dest='outputfile',

    14             help='Render output file. Type of output dependend on file extensions. Use png or jpg to render graph to image.'),

    15         make_option('--layout', '-l', action='store', dest='layout', default='dot',

    16             help='Layout to be used by GraphViz for visualization. Layouts: circo dot fdp neato nop nop1 nop2 twopi'),

    17     )

    18 

    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.")

    20     args = "[appname]"

    21     label = 'application name'

    22 

    23     requires_model_validation = True

    24     can_import_settings = True

    25 

    26     def handle(self, *args, **options):

    27         if len(args) < 1 and not options['all_applications']:

    28             raise CommandError("need one or more arguments for appname")

    29 

    30         dotdata = generate_dot(args, **options)

    31         if options['outputfile']:

    32             self.render_output(dotdata, **options)

    33         else:

    34             self.print_output(dotdata)

    35 

    36     def print_output(self, dotdata):

    37         print dotdata

    38 

    39     def render_output(self, dotdata, **kwargs):

    40         try:

    41             import pygraphviz

    42         except ImportError, e:

    43             raise CommandError("need pygraphviz python module ( apt-get install python-pygraphviz )")

    44 

    45         vizdata = ' '.join(dotdata.split("\n")).strip()

    46         version = pygraphviz.__version__.rstrip("-svn")

    47         try:

    48             if [int(v) for v in version.split('.')]<(0,36):

    49                 # HACK around old/broken AGraph before version 0.36 (ubuntu ships with this old version)

    50                 import tempfile

    51                 tmpfile = tempfile.NamedTemporaryFile()

    52                 tmpfile.write(vizdata)

    53                 tmpfile.seek(0)

    54                 vizdata = tmpfile.name

    55         except ValueError:

    56             pass

    57 

    58         graph = pygraphviz.AGraph(vizdata)

    59         graph.layout(prog=kwargs['layout'])

    60         graph.draw(kwargs['outputfile'])