From 9fa3d27c4bdd534eaff88ea2c4a7119e3174dbbf Mon Sep 17 00:00:00 2001
From: Brian Lee <brianklee@google.com>
Date: Wed, 6 Mar 2019 18:37:28 -0800
Subject: [PATCH] Add unicode string support to pyct pretty printer.
PiperOrigin-RevId: 237161597
---
tensorflow/python/autograph/pyct/pretty_printer.py | 9 +++++++--
tensorflow/python/autograph/pyct/pretty_printer_test.py | 9 +++++++++
2 files changed, 16 insertions(+), 2 deletions(-)
diff --git a/tensorflow/python/autograph/pyct/pretty_printer.py b/tensorflow/python/autograph/pyct/pretty_printer.py
index a92017f4142..9a4f509ec36 100644
--- a/tensorflow/python/autograph/pyct/pretty_printer.py
+++ b/tensorflow/python/autograph/pyct/pretty_printer.py
@@ -18,7 +18,9 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
+
import gast
+import six
import termcolor
@@ -106,9 +108,12 @@ class PrettyPrinter(gast.NodeVisitor):
self._print('%s%s=()' % (self._indent(), self._field(f)))
elif isinstance(v, gast.AST):
self.generic_visit(v, f)
- elif isinstance(v, str):
+ elif isinstance(v, six.binary_type):
self._print('%s%s=%s' % (self._indent(), self._field(f),
- self._value('"%s"' % v)))
+ self._value('b"%s"' % v)))
+ elif isinstance(v, six.text_type):
+ self._print('%s%s=%s' % (self._indent(), self._field(f),
+ self._value('u"%s"' % v)))
else:
self._print('%s%s=%s' % (self._indent(), self._field(f),
self._value(v)))
diff --git a/tensorflow/python/autograph/pyct/pretty_printer_test.py b/tensorflow/python/autograph/pyct/pretty_printer_test.py
index 1c76744547f..26d70f2e600 100644
--- a/tensorflow/python/autograph/pyct/pretty_printer_test.py
+++ b/tensorflow/python/autograph/pyct/pretty_printer_test.py
@@ -19,6 +19,7 @@ from __future__ import division
from __future__ import print_function
import ast
+import textwrap
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.platform import test
@@ -26,6 +27,14 @@ from tensorflow.python.platform import test
class PrettyPrinterTest(test.TestCase):
+ def test_unicode_bytes(self):
+ source = textwrap.dedent('''
+ def f():
+ return b'b', u'u', 'depends_py2_py3'
+ ''')
+ node = ast.parse(source)
+ self.assertIsNotNone(pretty_printer.fmt(node))
+
def test_format(self):
node = ast.FunctionDef(
name='f',