Support optional application-name argument, fix CamelCased objects in ForeignKeys
diff --git a/planetstack/dmdot b/planetstack/dmdot
index 2d95e9d..4353793 100755
--- a/planetstack/dmdot
+++ b/planetstack/dmdot
@@ -10,40 +10,61 @@
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "planetstack.settings")
from django.db.models.fields.related import ForeignKey
-from core.models import *
-try:
- output = sys.args[1]
-except:
- output = '-json'
+# defaults
+app = "core"
+output = "-json"
-g = globals()
+# 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:]:
+ models_module = getattr(models_module,part)
+
+PlCoreBase = getattr(models_module,"PlCoreBase")
+
model_classes = []
class_names = []
-for c in g.values():
+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.title() in class_names:
- print '\t"%s"->"%s";'%(c.__name__,f.name.title())
+ 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.title() in class_names:
+ 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(f.name.title())
+ d[c.__name__].append(linked_class.__name__)
except KeyError:
- d[c.__name__]=[f.name.title()]
+ d[c.__name__]=[linked_class.__name__]
print json.dumps(d,indent=4)