diff --git a/scripts/graph.py b/scripts/graph.py
index 8c74fcc..f588bf0 100644
--- a/scripts/graph.py
+++ b/scripts/graph.py
@@ -1,8 +1,13 @@
 import os, sys
+from os.path import isdir, join
+
+dir = sys.argv[1] if len(sys.argv) > 1 else '.'
 
 print 'digraph {'
-for i in os.listdir(sys.argv[1] if len(sys.argv) > 1 else '.'):
-    parts = i.rsplit('.', 1)[0].split('=')
+for i in os.listdir(dir):
+    if not isdir(join(dir, i)):
+        i = i.rsplit('.', 1)[0]
+    parts = i.split('=')
     print ' -> '.join(parts[:2])
     # TODO: add label (optional section 3)
 print '}'