From e6affeb79ee0cfda24b76368e3e788a7ec23df32 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 19 Mar 2018 14:09:52 -0700
Subject: [PATCH] Add a helper that allows constructing simple expression ASTs
 from string. Useful to simplify the representation of composite symbols, e.g.
 'py2tf.foo'.

PiperOrigin-RevId: 189638901
---
 tensorflow/contrib/py2tf/pyct/parser.py      | 22 ++++++++++++++++++++--
 tensorflow/contrib/py2tf/pyct/parser_test.py | 17 +++++++++++------
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/py2tf/pyct/parser.py b/tensorflow/contrib/py2tf/pyct/parser.py
index dc7df88..c961efa 100644
--- a/tensorflow/contrib/py2tf/pyct/parser.py
+++ b/tensorflow/contrib/py2tf/pyct/parser.py
@@ -29,12 +29,30 @@ from tensorflow.python.util import tf_inspect
 
 
 def parse_entity(entity):
-  """Return the AST of given entity."""
+  """Returns the AST of given entity."""
   source = tf_inspect.getsource(entity)
   source = textwrap.dedent(source)
   return parse_str(source), source
 
 
 def parse_str(src):
-  """Return the AST of given piece of code."""
+  """Returns the AST of given piece of code."""
   return gast.parse(src)
+
+
+def parse_expression(src):
+  """Returns the AST of given identifier.
+
+  Args:
+    src: A piece of code that represents a single Python expression
+  Returns:
+    A gast.AST object.
+  Raises:
+    ValueError: if src does not consist of a single Expression.
+  """
+  node = parse_str(src)
+  assert isinstance(node, gast.Module)
+  if len(node.body) != 1 and not isinstance(node.body[0], gast.Expr):
+    raise ValueError(
+        'Expected a single expression, found instead %s' % node.body)
+  return node.body[0].value
diff --git a/tensorflow/contrib/py2tf/pyct/parser_test.py b/tensorflow/contrib/py2tf/pyct/parser_test.py
index f35dfa0..c58ffc7 100644
--- a/tensorflow/contrib/py2tf/pyct/parser_test.py
+++ b/tensorflow/contrib/py2tf/pyct/parser_test.py
@@ -24,24 +24,29 @@ from tensorflow.contrib.py2tf.pyct import parser
 from tensorflow.python.platform import test
 
 
-def f(x):
-  return x + 1
-
-
 class ParserTest(test.TestCase):
 
   def test_parse_entity(self):
+
+    def f(x):
+      return x + 1
+
     mod, _ = parser.parse_entity(f)
     self.assertEqual('f', mod.body[0].name)
 
   def test_parse_str(self):
     mod = parser.parse_str(
         textwrap.dedent("""
-        def f(x):
-          return x + 1
+            def f(x):
+              return x + 1
     """))
     self.assertEqual('f', mod.body[0].name)
 
+  def test_parse_expression(self):
+    node = parser.parse_expression('a.b')
+    self.assertEqual('a', node.value.id)
+    self.assertEqual('b', node.attr)
+
 
 if __name__ == '__main__':
   test.main()
-- 
2.7.4