Handle comments, char refs, etc before and after root node.

kata198 · kata198 · commit 4eb09b1343a5 · 2016-07-25T12:38:10.000-04:00
diff --git a/AdvancedHTMLParser/Parser.py b/AdvancedHTMLParser/Parser.py
@@ -145,20 +145,26 @@ def handle_entityref(self, entity):
         '''
         if len(self.inTag) > 0:
             self.inTag[-1].appendText('&%s;' %(entity,))
+        else:
+            raise MultipleRootNodeException()
 
     def handle_charref(self, charRef):
         '''
             Internal for parsing
         '''
         if len(self.inTag) > 0:
             self.inTag[-1].appendText('&#%s;' %(charRef,))
+        else:
+            raise MultipleRootNodeException()
 
     def handle_comment(self, comment):
         '''
             Internal for parsing
         '''
         if len(self.inTag) > 0:
             self.inTag[-1].appendText('<!-- %s -->' %(comment,))
+        else:
+            raise MultipleRootNodeException()
 
     def handle_decl(self, decl):
         '''
diff --git a/tests/AdvancedHTMLParserTests/test_untaggedText.py b/tests/AdvancedHTMLParserTests/test_untaggedText.py
@@ -76,6 +76,51 @@ def test_textBeforeAndAfterRoot(self):
         assert strippedHTML.startswith('Hello') , 'Expected text before root tag to be retained, got "%s"' %(strippedHTML,)
         assert strippedHTML.endswith('World') , 'Expected text after root tag to be retained, got "%s"' %(strippedHTML,)
 
+    def test_commentRetained(self):
+        html = """<html>
+        <!-- CommentX -->
+        <body><span>Hello</span></body></html>"""
+
+        parser = AdvancedHTMLParser()
+        parser.parseStr(html)
+
+        retHTML = parser.getHTML()
+
+        assert 'CommentX' in retHTML, 'Expected to find comment, "CommentX" in returned HTML: "%s"' %(retHTML,)
+
+    def test_commentRetainedPriorRoot(self):
+        html = """<!-- CommentX --><html>
+        <body><span>Hello</span></body></html>"""
+
+        parser = AdvancedHTMLParser()
+        parser.parseStr(html)
+
+        retHTML = parser.getHTML()
+
+        assert 'CommentX' in retHTML, 'Expected to find comment, "CommentX" in returned HTML: "%s"' %(retHTML,)
+
+    def test_commentRetainedAfterRoot(self):
+        html = """<html>
+        <body><span>Hello</span></body></html><!-- CommentX -->"""
+
+        parser = AdvancedHTMLParser()
+        parser.parseStr(html)
+
+        retHTML = parser.getHTML()
+
+        assert 'CommentX' in retHTML, 'Expected to find comment, "CommentX" in returned HTML: "%s"' %(retHTML,)
+
+    def test_commentRetainedBeforeAndAfterRoot(self):
+        html = """<!-- CommentX --><html>
+        <body><span>Hello</span></body></html><!-- CommentY -->"""
+
+        parser = AdvancedHTMLParser()
+        parser.parseStr(html)
+
+        retHTML = parser.getHTML()
+
+        assert 'CommentX' in retHTML, 'Expected to find comment, "CommentX" in returned HTML: "%s"' %(retHTML,)
+        assert 'CommentY' in retHTML, 'Expected to find comment, "CommentY" in returned HTML: "%s"' %(retHTML,)
 
 
 
