Use six in richstring to be compatible with Python 3

2 files changed, 20 insertions, 2 deletions

diff --git a/requirements.txt b/requirements.txt
index 3d09f64..08a90f9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,4 @@
 unittest2
 nose
 pep8
+six
diff --git a/screenplain/richstring.py b/screenplain/richstring.py
index 43f5949..19fd45d 100644
--- a/screenplain/richstring.py
+++ b/screenplain/richstring.py
@@ -4,10 +4,21 @@
 
 import re
 import cgi
+import six
 
 _magic_re = re.compile(u'[\ue700-\ue705]')
 
 
+def _escape(s):
+    """Replaces special HTML characters like <
+    and non-ascii characters with ampersand escapes.
+
+    """
+    encoded = cgi.escape(s).encode('ascii', 'xmlcharrefreplace')
+    # In Py3, encoded is bytes type, so convert it to a string
+    return encoded.decode('ascii')
+
+
 class RichString(object):
     """A sequence of segments where each segment can have its own style."""
 
@@ -20,7 +31,10 @@ class RichString(object):
         return ' + '.join(repr(s) for s in self.segments)
 
     def __unicode__(self):
-        return ''.join(unicode(s) for s in self.segments)
+        return ''.join(six.text_type(s) for s in self.segments)
+
+    def __str__(self):
+        return self.__unicode__()
 
     def startswith(self, string):
         """Checks if the first segment in this string starts with a
@@ -89,6 +103,9 @@ class Segment(object):
     def __unicode__(self):
         return self.text
 
+    def __str__(self):
+        return self.text
+
     def __eq__(self, other):
         return (
             isinstance(other, Segment) and
@@ -109,7 +126,7 @@ class Segment(object):
         ordered_styles = self.get_ordered_styles()
         return (
             ''.join(style.start_html for style in ordered_styles) +
-            cgi.escape(self.text).encode('ascii', 'xmlcharrefreplace') +
+            _escape(self.text) +
             ''.join(style.end_html for style in reversed(ordered_styles))
         )