Add a utility that converts call keyword arguments into dicts, in AST space.

PiperOrigin-RevId: 190047495

diff --git a/tensorflow/contrib/py2tf/pyct/ast_util.py b/tensorflow/contrib/py2tf/pyct/ast_util.py
index f916775..6f7e656 100644 (file)
--- a/tensorflow/contrib/py2tf/pyct/ast_util.py
+++ b/tensorflow/contrib/py2tf/pyct/ast_util.py
@@ -94,3 +94,12 @@ def rename_symbols(node, name_map):
   elif isinstance(node, tuple):
     return tuple(renamer.visit(n) for n in node)
   return renamer.visit(node)
+
+
+def keywords_to_dict(keywords):
+  keys = []
+  values = []
+  for kw in keywords:
+    keys.append(gast.Str(kw.arg))
+    values.append(kw.value)
+  return gast.Dict(keys=keys, values=values)

diff --git a/tensorflow/contrib/py2tf/pyct/ast_util_test.py b/tensorflow/contrib/py2tf/pyct/ast_util_test.py
index a871cca..8d12367 100644 (file)
--- a/tensorflow/contrib/py2tf/pyct/ast_util_test.py
+++ b/tensorflow/contrib/py2tf/pyct/ast_util_test.py
@@ -21,6 +21,8 @@ from __future__ import print_function
 import ast
 
 from tensorflow.contrib.py2tf.pyct import ast_util
+from tensorflow.contrib.py2tf.pyct import compiler
+from tensorflow.contrib.py2tf.pyct import parser
 from tensorflow.contrib.py2tf.pyct import qual_names
 from tensorflow.python.platform import test
 
@@ -74,6 +76,17 @@ class AstUtilTest(test.TestCase):
     self.assertFalse(ret is new_node.body[0])
     self.assertFalse(hasattr(new_node.body[0], '__foo'))
 
+  def test_keywords_to_dict(self):
+    keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
+    d = ast_util.keywords_to_dict(keywords)
+    # Make sure we generate a usable dict node by attaching it to a variable and
+    # compiling everything.
+    output = parser.parse_str('b = 3')
+    output.body += (ast.Assign([ast.Name(id='d', ctx=ast.Store())], d),)
+    result, _ = compiler.ast_to_object(output)
+    self.assertDictEqual(result.d, {'a': 3, 'c': 1, 'd': 'e'})
+    print(d)
+
 
 if __name__ == '__main__':
   test.main()