Merge branch 'add-styles'

paetzke · Mar 8, 2014 · 816ac08 · 816ac08
2 parents 195ee5c + 7030c5b
commit 816ac08
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 5 deletions.
diff --git a/dotter/__init__.py b/dotter/__init__.py
@@ -6,7 +6,7 @@
 All rights reserved.
 
 """
-from .dotter import Dotter, Program, RankType, Shape
+from .dotter import Dotter, Program, RankType, Shape, Style
 
 __version__ = '0.3.0'
 __author__ = 'Friedrich Paetzke'

diff --git a/dotter/dotter.py b/dotter/dotter.py
@@ -103,6 +103,23 @@ class Shape:
     Utr = 'utr'
 
 
+class Style:
+
+    """
+    These values can be used as valid node style values.
+    See http://www.graphviz.org/doc/info/shapes.html#d:style for more info.
+
+    """
+    Bold = 'bold'
+    Dashed = 'dashed'
+    Diagonals = 'diagonals'
+    Dotted = 'dotted'
+    Filled = 'filled'
+    Invisible = 'invisible'
+    Rounded = 'rounded'
+    Solid = 'solid'
+
+
 class Dotter:
 
     def __init__(self, directed=True, output_to_file=True,
@@ -170,12 +187,14 @@ def add_edge(self, node1, node2, label=None):
             fmt += ' [label="{0}"]'.format(label)
         self.execute(fmt.format(Dotter.escape(node1), Dotter.escape(node2)))
 
-    def add_node(self, node, font=None, fontsize=None, label=None, shape=None, url=None):
+    def add_node(self, node, font=None, fontsize=None, label=None, shape=None,
+                 url=None, styles=None):
         self.execute('{0}'.format(Dotter.escape(node)))
         self.set_label(node, label if label else node)
-        self.node_attributes(node, font, fontsize, shape, url)
+        self.node_attributes(node, font, fontsize, shape, url, styles)
 
-    def node_attributes(self, node, font=None, fontsize=None, shape=None, url=None):
+    def node_attributes(self, node, font=None, fontsize=None, shape=None,
+                        url=None, styles=None):
         if font:
             cmd = '{0} [fontname="{1}"]'.format(Dotter.escape(node), font)
             self.execute(cmd)
@@ -188,6 +207,10 @@ def node_attributes(self, node, font=None, fontsize=None, shape=None, url=None):
         if url:
             cmd = '{0} [URL="{1}"]'.format(Dotter.escape(node), url)
             self.execute(cmd)
+        if styles:
+            style = ', '.join(styles)
+            cmd = '{0} [style="{1}"]'.format(Dotter.escape(node), style)
+            self.execute(cmd)
 
     def set_label(self, node, label):
         self.execute('{0} [label="{1}"]'.format(Dotter.escape(node), label))

diff --git a/tests/test_dotter.py b/tests/test_dotter.py
@@ -8,9 +8,10 @@
 """
 import os
 
-from dotter import Dotter, Program, RankType, Shape
 from pytest import raises
 
+from dotter import Dotter, Program, RankType, Shape, Style
+
 
 def load_data(filename):
     test_data = os.path.dirname(os.path.abspath(os.path.dirname(__file__)))
@@ -48,6 +49,18 @@ def test_addnode_label():
     assert expected == dotter.commands
 
 
+def test_addnode_styles():
+    dotter = Dotter()
+    dotter.add_node('a', styles=["diagonals", "filled", Style.Bold])
+    expected = ['digraph',
+                ' {',
+                'gb',
+                'gb [label="a"]',
+                'gb [style="diagonals, filled, bold"]',
+                ]
+    assert expected == dotter.commands
+
+
 def test_set_position():
     dotter = Dotter(program=Program.Neato)
     dotter.add_node('a')