Make our enhanced fromstring function accept bytes as well

Fixes: #21

Author: frenzymadness

lxml_html_clean/clean.py

@@ -90,15 +90,19 @@ def _has_javascript_scheme(s):
 # - 0A - Line Feed
 # - 0B - Vertical tab
 # - 0D - Carriage Return
-_ascii_control_characters = re.compile(r"[\x00-\x08\x0C\x0E-\x1F\x7F]")
+_ascii_control_characters_str = re.compile("[\x00-\x08\x0C\x0E-\x1F\x7F]")
+_ascii_control_characters_bytes = re.compile(b"[\x00-\x08\x0C\x0E-\x1F\x7F]")
 
 
-def fromstring(string):
+def fromstring(data):
     """
     Enhanced fromstring function that removes ASCII control chars
     before passing the input to the original lxml.html.fromstring.
     """
-    return lxml_fromstring(_ascii_control_characters.sub("", string))
+    if isinstance(data, bytes):
+        return lxml_fromstring(_ascii_control_characters_bytes.sub(b"", data))
+    else:
+        return lxml_fromstring(_ascii_control_characters_str.sub("", data))
 
 
 # This regular expression is inspired by the one in urllib3.

tests/test_clean.py

@@ -355,6 +355,12 @@ def test_ascii_control_chars_removed(self):
         cleaner = Cleaner()
         self.assertEqual(expected, cleaner.clean_html(html))
 
+    def test_ascii_control_chars_removed_from_bytes(self):
+        html = b"""<a href="java\x1bscript:alert()">Link</a>"""
+        expected = b"""<a href="">Link</a>"""
+        cleaner = Cleaner()
+        self.assertEqual(expected, cleaner.clean_html(html))
+
     def test_memory_usage_many_elements_with_long_tails(self):
         comment = "<!-- foo bar baz -->\n"
         empty_line = "\t" * 10 + "\n"