Merge pull request #1 from eeshugerman/master

highlight and label release revisions

README.md

@@ -1,7 +1,7 @@
-####Installation
+#### Installation
 `pip install alembic-viz --user`
 
-####Usage
+#### Usage
 ```bash
 Usage: alembic-viz [OPTIONS]
 
@@ -13,8 +13,8 @@ Options:
   --help                  Show this message and exit.
 ```
 
-####Todo
+#### Todo
 * add revision commit messages to graph
 * handle `depends_on` relationships properly
 * add test cases
-* ...
+* ...

alembic_viz/alembic_viz.py

@@ -21,7 +21,7 @@ def cli(config, name, filename, format):
         revisions = get_revisions(alembic_config)
     except CommandError as e:
         sys.exit(e)
-    dot = generate_revision_graph(revisions, format)
+    dot = generate_revision_graph(revisions, format, alembic_config)
     dot.render(filename=filename)
 
 

alembic_viz/utils.py

@@ -3,6 +3,21 @@
 from graphviz import Digraph
 
 
+RELEASE_REVISION_TAG = '[RELEASE-MERGE]'
+
+
+def is_release_revision(revision):
+    rval = revision.doc.startswith(RELEASE_REVISION_TAG)
+    return rval
+
+
+def get_node_name(revision):
+    if is_release_revision(revision):
+        return revision.doc[len(RELEASE_REVISION_TAG) + 1:]
+    else:
+        return revision.revision
+
+
 def get_revisions(config, rev_range=None):
     script = ScriptDirectory.from_config(config)
 
@@ -19,16 +34,27 @@ def get_revisions(config, rev_range=None):
     )
 
 
-def generate_revision_graph(revisions, fmt):
+def generate_revision_graph(revisions, fmt, config):
+    script = ScriptDirectory.from_config(config)
     dot = Digraph(format=fmt)
     for revision in revisions:
-        dot.node(revision.revision)
+        color = 'turquoise' if is_release_revision(revision) else None
+        node_name = get_node_name(revision)
+        dot.node(node_name, color=color)
+
         if revision.down_revision is None:
-            dot.edge('base', revision.revision)
+            dot.edge('base', get_node_name(revision))
             continue
+
         if isinstance(revision.down_revision, str):
-            dot.edge(revision.down_revision, revision.revision)
+            down_revision = script.get_revision(revision.down_revision)
+            down_node_name = get_node_name(down_revision)
+            dot.edge(down_node_name, get_node_name(revision))
             continue
+
         for down_revision in revision.down_revision:
-            dot.edge(down_revision, revision.revision)
+            down_revision = script.get_revision(down_revision)
+            down_node_name = get_node_name(down_revision)
+            dot.edge(down_node_name, node_name)
+
     return dot