Start work on getting text handled properly on PY3

1 files changed, 27 insertions, 7 deletions

diff --git a/python/google/protobuf/text_format.py b/python/google/protobuf/text_format.py
index 950bec16..7bdde33f 100755
--- a/python/google/protobuf/text_format.py
+++ b/python/google/protobuf/text_format.py
@@ -67,6 +67,29 @@ class Error(Exception):
 class ParseError(Error):
   """Thrown in case of ASCII parsing error."""
 
+class TextWriter(object):
+  def __init__(self, as_utf8):
+    self._utf8 = as_utf8
+    if as_utf8:
+      self._writer = io.BytesIO()
+    else:
+      self._writer = io.StringIO()
+
+  def write(self, val):
+    if self._utf8:
+      if isinstance(val, six.text_type):
+        val = val.encode('utf-8')
+    else:
+      if isinstance(val, bytes):
+        val = val.decode('utf-8')
+    return self._writer.write(val)
+
+  def close(self):
+    return self._writer.close()
+
+  def getvalue(self):
+    return self._writer.getvalue()
+
 
 def MessageToString(message, as_utf8=False, as_one_line=False,
                     pointy_brackets=False, use_index_order=False,
@@ -92,7 +115,7 @@ def MessageToString(message, as_utf8=False, as_one_line=False,
   Returns:
     A string of the text formatted protocol buffer message.
   """
-  out = io.BytesIO()
+  out = TextWriter(as_utf8)
   PrintMessage(message, out, as_utf8=as_utf8, as_one_line=as_one_line,
                pointy_brackets=pointy_brackets,
                use_index_order=use_index_order,
@@ -159,11 +182,7 @@ def PrintField(field, value, out, indent=0, as_utf8=False, as_one_line=False,
     # For groups, use the capitalized name.
     out.write(field.message_type.name)
   else:
-    if isinstance(field.name, six.text_type):
-      name = field.name.encode('utf-8')
-    else:
-      name = field.name
-    out.write(name)
+    out.write(field.name)
 
   if field.cpp_type != descriptor.FieldDescriptor.CPPTYPE_MESSAGE:
     # The colon is optional in this case, but our cross-language golden files
@@ -226,7 +245,8 @@ def PrintFieldValue(field, value, out, indent=0, as_utf8=False,
       out_as_utf8 = False
     else:
       out_as_utf8 = as_utf8
-    out.write(text_encoding.CEscape(out_value, out_as_utf8))
+    out_text = text_encoding.CEscape(out_value, out_as_utf8)
+    out.write(out_text)
     out.write('\"')
   elif field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_BOOL:
     if value: