From 30931bfe57f6e86d9c0422b4982c5bac8effbd0c Mon Sep 17 00:00:00 2001
From: James Scott-Brown <james@jamesscottbrown.com>
Date: Sat, 18 Jan 2020 11:19:47 +0000
Subject: [PATCH] Pretty print (#26)

* add support for pretty-printing

* bump release number
---
 README.md         |  9 ++++++++-
 pyyed/__init__.py | 14 +++++++++++---
 setup.py          |  2 +-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index fa43e8a..1e60095 100644
--- a/README.md
+++ b/README.md
@@ -32,11 +32,18 @@ g.add_node('foobar', label="""Multi
 g.add_edge('foo', 'foo1', label="EDGE!", width="3.0", color="#0000FF", 
                arrowhead="white_diamond", arrowfoot="standard", line_type="dotted")
 
-print g.get_graph()
+print(g.get_graph())
 
 # To write to file:
 with open('test_graph.graphml', 'w') as fp:
     fp.write(g.get_graph())
+
+# Or:
+g.write_graph('example.graphml')
+
+# Or, to pretty-print with whitespace:
+g.write_graph('pretty_example.graphml', pretty_print=True)
+
 ```
 
 Saving this to a file with a ``.graphml`` extension, opening in yEd, applying  ``Tools -> Fit Node to Label`` and ``Layout -> One-click layout`` produces something like the following:
diff --git a/pyyed/__init__.py b/pyyed/__init__.py
index b52b3dd..61c35ae 100644
--- a/pyyed/__init__.py
+++ b/pyyed/__init__.py
@@ -1,5 +1,6 @@
 import sys
 import xml.etree.ElementTree as ET
+from xml.dom import minidom
 
 node_shapes = ["rectangle", "rectangle3d", "roundrectangle", "diamond", "ellipse",
                "fatarrow", "fatarrow2", "hexagon", "octagon", "parallelogram",
@@ -333,10 +334,17 @@ def construct_graphml(self):
 
         self.graphml = graphml
 
-    def write_graph(self, filename):
+    def write_graph(self, filename, pretty_print=False):
         self.construct_graphml()
-        tree = ET.ElementTree(self.graphml)
-        tree.write(filename)
+
+        if pretty_print:
+            raw_str = ET.tostring(self.graphml)
+            pretty_str = minidom.parseString(raw_str).toprettyxml()
+            with open(filename, 'w') as f:
+                f.write(pretty_str)
+        else:
+            tree = ET.ElementTree(self.graphml)
+            tree.write(filename)
 
     def get_graph(self):
         self.construct_graphml()
diff --git a/setup.py b/setup.py
index 16fb0f6..6141a62 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@
 
 
 setup(name='pyyed',
-      version='1.2',
+      version='1.3',
       description='A simple Python library to export graphs to the yEd graph editor',
 
       author='James Scott-Brown',