good good, all tests pass except the one that needs to go

Marc DM · Marc DM · commit 07bbab1a3f4b · 2013-06-18T18:15:03.000-05:00
diff --git a/html5lib/serializer/htmlserializer.py b/html5lib/serializer/htmlserializer.py
@@ -1,5 +1,5 @@
 from __future__ import absolute_import, division, unicode_literals
-from six import string_types
+from six import text_type
 
 import gettext
 _ = gettext.gettext
@@ -154,14 +154,14 @@ def __init__(self, **kwargs):
         self.strict = False
 
     def encode(self, string):
-        assert(isinstance(string, string_types))
+        assert(isinstance(string, text_type))
         if self.encoding:
             return string.encode(self.encoding, unicode_encode_errors)
         else:
             return string
 
     def encodeStrict(self, string):
-        assert(isinstance(string, string_types))
+        assert(isinstance(string, text_type))
         if self.encoding:
             return string.encode(self.encoding, "strict")
         else:
diff --git a/html5lib/tests/test_six_encoding.py b/html5lib/tests/test_six_encoding.py
@@ -1,18 +1,20 @@
 
-from html5lib import html5parser, treewalkers, serializer
-from nose.tools import eq_
+from html5lib import html5parser, treewalkers, treebuilders, serializer
 
 
-def test_treewalker6():
+def test_treewalker_six_mix():
     """Str/Unicode mix. If str attrs added to tree"""
 
     text = '<a href="http://example.com">Example</a>'
-    end_text = '<a href="http://example.com" class="test123">Example</a>'
-    parser = html5parser.HTMLParser()
-    walker = treewalkers.getTreeWalker('etree')
+    end_texts = ('<a href="http://example.com" class="test123">Example</a>',
+                 '<a class="test123" href="http://example.com">Example</a>')
+    parser = html5parser.HTMLParser(tree=treebuilders.getTreeBuilder('dom'))
+    walker = treewalkers.getTreeWalker('dom')
     serializr = serializer.HTMLSerializer(quote_attr_values=True)
     domtree = parser.parseFragment(text)
 
     # at this point domtree should be a DOCUMENT_FRAGMENT
     domtree[0].set('class', 'test123')
-    eq_(end_text, serializr.render(walker(domtree)))
+    out = serializr.render(walker(domtree))
+    if not out in end_texts:
+        raise AssertionError('%r not in %r' % (out, end_texts))
diff --git a/html5lib/tests/test_treewalkers.py b/html5lib/tests/test_treewalkers.py
@@ -310,3 +310,41 @@ def test_treewalker():
                                                                "document")]
                 errors = errors.split("\n")
                 yield runTreewalkerTest, innerHTML, input, expected, errors, treeCls
+
+
+def set_attribute_on_first_child(docfrag, name, value, treeName):
+    """naively sets an attribute on the first child of the document
+    fragment passed in"""
+    setter = {'ElementTree': lambda d: d[0].set,
+              'DOM': lambda d: d.firstChild.setAttribute}
+    setter['PullDOM'] = setter['DOM']
+    try:
+        setter.get(treeName, setter['ElementTree'])(docfrag)(name, value)
+    except TypeError:
+        setter['DOM'](docfrag)(name, value)
+
+
+def runTreewalkerEditTest(intext, expected, attrs_to_add, tree):
+    """tests what happens when we add attributes to the intext"""
+    treeName, treeClass = tree
+    parser = html5parser.HTMLParser(tree=treeClass["builder"])
+    document = parser.parseFragment(intext)
+    for nom, val in attrs_to_add:
+        set_attribute_on_first_child(document, nom, val, treeName)
+
+    document = treeClass.get("adapter", lambda x: x)(document)
+    output = convertTokens(treeClass["walker"](document))
+    output = attrlist.sub(sortattrs, output)
+    if not output in expected:
+        raise AssertionError('%r not in %r' % (output, expected))
+
+
+def test_treewalker_six_mix():
+    """Str/Unicode mix. If str attrs added to tree"""
+
+    intext = '<a href="http://example.com">Example</a>'
+    expected = '<a>\n  class="test123"\n  href="http://example.com"\n  "Example"'
+    attrs = [('class', 'test123')]
+
+    for tree in treeTypes.items():
+        yield runTreewalkerEditTest, intext, expected, attrs, tree
diff --git a/html5lib/treewalkers/_base.py b/html5lib/treewalkers/_base.py
@@ -8,6 +8,24 @@
 spaceCharacters = "".join(spaceCharacters)
 
 
+def to_text(s, blank_if_none=True):
+    """Wrapper around six.text_type to convert None to empty string"""
+    if s is None:
+        if blank_if_none:
+            return ""
+        else:
+            return None
+    elif isinstance(s, text_type):
+        return s
+    else:
+        return text_type(s)
+
+
+def is_text_or_none(string):
+    """Wrapper around isinstance(string_types) or is None"""
+    return string is None or isinstance(string, string_types)
+
+
 class TreeWalker(object):
     def __init__(self, tree):
         self.tree = tree
@@ -26,8 +44,8 @@ def emptyTag(self, namespace, name, attrs, hasChildren=False):
                    isinstance(value, string_types)
                    for (namespace, name), value in attrs.items())
 
-        yield {"type": "EmptyTag", "name": text_type(name),
-               "namespace": text_type(namespace),
+        yield {"type": "EmptyTag", "name": to_text(name, False),
+               "namespace": to_text(namespace),
                "data": attrs}
         if hasChildren:
             yield self.error(_("Void element has children"))
@@ -42,22 +60,24 @@ def startTag(self, namespace, name, attrs):
 
         return {"type": "StartTag",
                 "name": text_type(name),
-                "namespace": text_type(namespace),
-                "data": attrs}
+                "namespace": to_text(namespace),
+                "data": dict(((to_text(ns), to_text(na, False)),
+                              to_text(va, False))
+                             for (ns, na), va in attrs.items())}
 
     def endTag(self, namespace, name):
         assert namespace is None or isinstance(namespace, string_types), type(namespace)
         assert isinstance(name, string_types), type(namespace)
 
         return {"type": "EndTag",
-                "name": text_type(name),
-                "namespace": text_type(namespace),
+                "name": to_text(name, False),
+                "namespace": to_text(namespace),
                 "data": {}}
 
     def text(self, data):
         assert isinstance(data, string_types), type(data)
 
-        data = data
+        data = to_text(data)
         middle = data.lstrip(spaceCharacters)
         left = data[:len(data) - len(middle)]
         if left:
@@ -73,23 +93,23 @@ def text(self, data):
     def comment(self, data):
         assert isinstance(data, string_types), type(data)
 
-        return {"type": "Comment", "data": data}
+        return {"type": "Comment", "data": text_type(data)}
 
     def doctype(self, name, publicId=None, systemId=None, correct=True):
-        assert name is None or isinstance(name, string_types), type(name)
-        assert publicId is None or isinstance(publicId, string_types), type(publicId)
-        assert systemId is None or isinstance(systemId, string_types), type(systemId)
+        assert is_text_or_none(name), type(name)
+        assert is_text_or_none(publicId), type(publicId)
+        assert is_text_or_none(systemId), type(systemId)
 
         return {"type": "Doctype",
-                "name": name if name is not None else "",
-                "publicId": publicId,
-                "systemId": systemId,
-                "correct": correct}
+                "name": to_text(name),
+                "publicId": to_text(publicId),
+                "systemId": to_text(systemId),
+                "correct": to_text(correct)}
 
     def entity(self, name):
         assert isinstance(name, string_types), type(name)
 
-        return {"type": "Entity", "name": name}
+        return {"type": "Entity", "name": text_type(name)}
 
     def unknown(self, nodeType):
         return self.error(_("Unknown node type: ") + nodeType)
