Fix pretty-printing of comments and PIs when written before or after …

…the document root element

streamxmlwriter.py

@@ -323,15 +323,25 @@ def declaration(self):
             self._wrote_declaration = True
     xml = declaration
 
+    def _comment_or_pi(self, data):
+        """Write a comment or PI, using special rules for
+        pretty-printing."""
+        self._close_start()
+        if self._pretty_print:
+            if ((self._tags and not self._wrote_data) or
+                (self._started and not self._tags)):
+                self.write("\n" + INDENT * len(self._tags))
+        self.write(data)
+        if self._pretty_print and not self._started:
+            self.write("\n")
+
     def comment(self, data):
         """Add an XML comment."""
-        self._close_start()
-        self.write("<!--" + escape_cdata(data, self.encoding) + "-->")
+        self._comment_or_pi("<!--" + escape_cdata(data, self.encoding) + "-->")
 
     def pi(self, target, data):
         """Add an XML processing instruction."""
-        self._close_start()
-        self.write("<?" + target + " " + data + "?>")
+        self._comment_or_pi("<?" + target + " " + data + "?>")
 
     def close(self):
         """Close all open elements."""

test/test_streamxmlwriter.py

@@ -95,7 +95,22 @@ def testClose(self):
         w.close()
         self.assertEqual(out.getvalue(), "<a><b /></a>")
 
-    def testPrettyPrint(self):
+    def testDeclarationLateDeclarationRaisesSyntaxError(self):
+        w, out = writer_and_output()
+        w.start("a")
+        self.assertRaises(XMLSyntaxError, w.declaration)
+
+    def testIgnoreDoubleDeclaration(self):
+        w, out = writer_and_output()
+        w.declaration()
+        w.declaration()
+        w.close()
+        self.assertEqual(out.getvalue(),
+                         "<?xml version='1.0' encoding='utf-8'?>")
+
+
+class PrettyPrintTestCase(unittest.TestCase):
+    def testSimple(self):
         w, out = writer_and_output(pretty_print=True)
         w.start("a")
         w.start("b")
@@ -109,18 +124,58 @@ def testPrettyPrint(self):
         w.close()
         self.assertEqual(out.getvalue(), "<a>\n  <b>foo</b>\n  <b>bar</b>\n  <b>\n    <c />\n  </b>\n</a>")
 
-    def testDeclarationLateDeclarationRaisesSyntaxError(self):
-        w, out = writer_and_output()
+    def testComment(self):
+        w, out = writer_and_output(pretty_print=True)
         w.start("a")
-        self.assertRaises(XMLSyntaxError, w.declaration)
+        w.comment("comment")
+        w.start("b")
+        w.close()
+        self.assertEqual(out.getvalue(),
+                         "<a>\n  <!--comment-->\n  <b />\n</a>")
 
-    def testIgnoreDoubleDeclaration(self):
-        w, out = writer_and_output()
-        w.declaration()
-        w.declaration()
+    def testCommentBeforeRoot(self):
+        w, out = writer_and_output(pretty_print=True)
+        w.comment("comment")
+        w.start("a")
         w.close()
         self.assertEqual(out.getvalue(),
-                         "<?xml version='1.0' encoding='utf-8'?>")
+                         "<!--comment-->\n<a />")
+
+    def testCommentAfterRoot(self):
+        w, out = writer_and_output(pretty_print=True)
+        w.start("a")
+        w.end()
+        w.comment("comment")
+        w.close()
+        self.assertEqual(out.getvalue(),
+                         "<a />\n<!--comment-->")
+
+    def testPI(self):
+        w, out = writer_and_output(pretty_print=True)
+        w.start("a")
+        w.pi("foo", "bar")
+        w.start("b")
+        w.close()
+        self.assertEqual(out.getvalue(),
+                         "<a>\n  <?foo bar?>\n  <b />\n</a>")
+
+    def testPIBeforeRoot(self):
+        w, out = writer_and_output(pretty_print=True)
+        w.pi("foo", "bar")
+        w.start("a")
+        w.close()
+        self.assertEqual(out.getvalue(),
+                         "<?foo bar?>\n<a />")
+
+    def testPIAfterRoot(self):
+        w, out = writer_and_output(pretty_print=True)
+        w.start("a")
+        w.end()
+        w.pi("foo", "bar")
+        w.close()
+        self.assertEqual(out.getvalue(),
+                         "<a />\n<?foo bar?>")
+
 
 class NamespaceTestCase(unittest.TestCase):
     def testSimple(self):