Add 'body' and 'head' properties to AdvancedHTMLParser

Author: kata198

AdvancedHTMLParser/Parser.py

@@ -419,6 +419,14 @@ def getFirstElementCustomFilter(self, filterFunc, root='root'):
 
         return None
 
+    @property
+    def body(self):
+        return self.getFirstElementCustomFilter(lambda em : em.tagName == 'body')
+
+    @property
+    def head(self):
+        return self.getFirstElementCustomFilter(lambda em : em.tagName == 'head')
+
     def contains(self, em):
         '''
             Checks if #em is found anywhere within this element tree

tests/AdvancedHTMLParserTests/test_General.py

@@ -178,6 +178,24 @@ def test_escapeQuotes(self):
 
         assert 'onclick="alert("hi");"' in tag.outerHTML , 'Expected to escape quotes in attribute. Got: %s' %(tag.outerHTML, )
 
+    def test_getBody(self):
+        parser = AdvancedHTMLParser.AdvancedHTMLParser()
+        parser.parseStr(self.testHTML)
+
+        bodyEm = parser.body
+
+        assert bodyEm , 'Expected ".body" property to fetch an element'
+        assert bodyEm.tagName == 'body' , 'Got wrong body element'
+
+    def test_getHead(self):
+        parser = AdvancedHTMLParser.AdvancedHTMLParser()
+        parser.parseStr(self.testHTML)
+
+        headEm = parser.head
+
+        assert headEm , 'Expected ".head" property to fetch an element'
+        assert headEm.tagName == 'head' , 'Got wrong head element'
+
 
 if __name__ == '__main__':
     subprocess.Popen('GoodTests.py "%s"' %(sys.argv[0],), shell=True).wait()