planetstack/dmdot

#!/usr/bin/python

import os
import pdb
import sys
import json

sys.path.append('.')

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "planetstack.settings")

from django.db.models.fields.related import ForeignKey

# defaults
app = "core"
output = "-json"

# syntax: dmdot [-json | -dot] [app_name]

# poor man's argument parser
for arg in sys.argv[1:]:
    if arg.startswith("-"):
        output = arg
    else:
        app = arg

app = app + ".models"
#models_module = imp.load_source(app, ".")
models_module = __import__(app)
for part in app.split(".")[1:]:
    if hasattr(models_module, "PlCoreBase"):
        break
    models_module = getattr(models_module,part)

PlCoreBase = getattr(models_module,"PlCoreBase")

model_classes = []
class_names = []
lower_class_names = {}
for classname in dir(models_module):
        c = getattr(models_module, classname, None)
	if type(c)==type(PlCoreBase):
		model_classes.append(c)
		class_names.append(c.__name__)
                lower_class_names[c.__name__.lower()] = c

# django doesn't use the correct case in field.name.title() for objects that
# have CamelCased class names. So, compare everything in lower case.

if (output=='-dot'):
	print "digraph plstack {";
	for c in model_classes:
		fields = c._meta.fields
		for f in fields:
			if type(f)==ForeignKey and f.name.lower() in lower_class_names:
                                linked_class = lower_class_names[f.name.lower()]
				print '\t"%s"->"%s";'%(c.__name__,linked_class.__name__)
	print "}\n";
elif (output=='-json'):
	d = {}
	for c in model_classes:
		fields = c._meta.fields
		for f in fields:
			if type(f)==ForeignKey and f.name.lower() in lower_class_names:
                                linked_class = lower_class_names[f.name.lower()]
				try:
					d[c.__name__].append(linked_class.__name__)
				except KeyError:
					d[c.__name__]=[linked_class.__name__]
	print json.dumps(d,indent=4)