Management Command to Visualize Model Relationships

# -*- coding: utf-8 -*-

"""
Django management command to visualize model relations
management/commands/view_models_relation.py

./manage.py view_models_relation

+ django.contrib.auth.models.Group
+ django.contrib.auth.models.User
| + django.contrib.auth.models.Message
| + django.contrib.admin.models.LogEntry
+ django.contrib.contenttypes.models.ContentType
| + django.contrib.auth.models.Permission
| + django.contrib.admin.models.LogEntry
+ django.contrib.sessions.models.Session
+ django.contrib.sites.models.Site
+ south.models.MigrationHistory
"""

from optparse import make_option
from django.core.management.base import BaseCommand
from django.db import models
from django.utils.encoding import smart_str
from django.conf import settings
from django.utils.importlib import import_module

class Command(BaseCommand):
    help = 'print model relation'
    args = '<app_name>'
    
    option_list = BaseCommand.option_list + (
        make_option('--verbose',
            action='store_true',
            dest='verbose',
            default=False,
            help='Verbose message output'),
    )
    
    def echo(self, message):
        """
        Example: self.echo(self.style.NOTICE('NG!'))
        """
        self.stdout.write(smart_str(message, errors='ignore'))
        self.stdout.write('\n')
        self.stdout.flush()
    
    def echo_verbose(self, message):
        if self._verbose:
            return self.echo(message)
    
    def handle(self, app_label=None, **options):
        
        self._verbose=options.get('verbose', False)
        
        # Import APP
        for app_name in settings.INSTALLED_APPS:
            try:
                import_module('.management', app_name)
            except ImportError as exc:
                msg = exc.args[0]
                if not msg.startswith('No module named') or 'management' not in msg:
                    raise
        
        if app_label:
            app_list = [models.get_app(app_label)]
        else:
            app_list = models.get_apps()
        
        self.childs_dict = {} # List of children
        #self.parents_dict = {} # List of parents
        primitive_models = [] # No relations at all
        for app in app_list:
            if not app_label:
                self.echo_verbose('=' * 50)
                self.echo_verbose('App: %s' % app)
            
            app_models = models.get_models(app)
            
            if not app_models:
                self.echo_verbose('No app models.')
            
            for app_model in app_models:
                self.echo_verbose('-' * 50)
                app_model_name = self.get_model_path(app_model)
                self.echo_verbose('app_model_name=%s' % app_model_name)
                has_parent = False
                for field in app_model._meta.fields:
                    field_class_name = field.__class__.__name__
                    if field_class_name == 'ForeignKey' or \
                        field_class_name == 'OneToOneField' or \
                        field_class_name == 'ManyToManyField':
                        self.echo_verbose('related_field: %s' % field.name)
                        related_model = field.related.parent_model
                        rmn = self.get_model_path(related_model)
                        self.echo_verbose('related_model: %s' % rmn)
                        self.append_dict_list(self.childs_dict, rmn, app_model_name)
                        #self.append_dict_list(self.parents_dict, app_model_name, rmn)
                        has_parent = True
                if not has_parent:
                    primitive_models.append(app_model_name)
        
        for primitive_model in primitive_models:
            self.print_child(primitive_model)
    
    def print_child(self, model_name, depth=0):
        indent = ("|  " * depth)
        self.echo(indent +"+ "+ model_name)
        if model_name in self.childs_dict:
            child_models = self.childs_dict[model_name]
            for child_model in child_models:
                if child_model == model_name:
                    self.echo(indent +"+ " + child_model + " (self relation)")
                else:
                    self.print_child(child_model, depth=depth+1)
    
    def append_dict_list(self, d, k, v):
        """
        Dictionary d has a list as its value.
        Add an item to that list. If it doesn't exist, create a new list.
        """
        if k in d:
            d[k].append(v)
        else:
            d[k] = [v]

    def get_model_path(self, model_class):
        return "%s.%s" % (model_class.__module__, model_class.__name__)