More AST utils. Includes support for persistent annotations and copying the multiple assignment walking code out of transformer.py (to be removed later).

Persistent annotations require:
 * allow copy_clean to keep user-specified annotations
 * ensure rename_symbols does not destroy annotations

PiperOrigin-RevId: 204336881

2 files changed, 188 insertions, 78 deletions

diff --git a/tensorflow/contrib/autograph/pyct/ast_util.py b/tensorflow/contrib/autograph/pyct/ast_util.py
index c4f82d1170..0cf87dd8d3 100644
--- a/tensorflow/contrib/autograph/pyct/ast_util.py
+++ b/tensorflow/contrib/autograph/pyct/ast_util.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Copy an AST tree, discarding annotations."""
+"""AST manipulation utilities."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -26,47 +26,53 @@ from tensorflow.contrib.autograph.pyct import anno
 from tensorflow.contrib.autograph.pyct import parser
 
 
-class CleanCopier(gast.NodeVisitor):
-  """Copies AST nodes.
+class CleanCopier(object):
+  """NodeTransformer-like visitor that copies an AST."""
 
-  The copied nodes will ignore almost all fields that are prefixed by '__'.
-  Exceptions make some annotations.
-  """
+  def __init__(self, preserve_annos):
+    super(CleanCopier, self).__init__()
+    self.preserve_annos = preserve_annos
 
-  # TODO(mdan): Parametrize which annotations get carried over.
+  def copy(self, node):
+    """Returns a deep copy of node (excluding some fields, see copy_clean)."""
+
+    if isinstance(node, list):
+      return [self.copy(n) for n in node]
+    elif isinstance(node, tuple):
+      return tuple(self.copy(n) for n in node)
+    elif not isinstance(node, (gast.AST, ast.AST)):
+      # Assuming everything that's not an AST, list or tuple is a value type
+      # and may simply be assigned.
+      return node
+
+    assert isinstance(node, (gast.AST, ast.AST))
 
-  def generic_visit(self, node):
     new_fields = {}
     for f in node._fields:
-      if f.startswith('__'):
-        continue
-      if not hasattr(node, f):
-        continue
-      v = getattr(node, f)
-      if isinstance(v, list):
-        v = [self.generic_visit(n) for n in v]
-      elif isinstance(v, tuple):
-        v = tuple(self.generic_visit(n) for n in v)
-      elif isinstance(v, (gast.AST, ast.AST)):
-        v = self.generic_visit(v)
-      else:
-        # Assume everything else is a value type.
-        pass
-      new_fields[f] = v
+      if not f.startswith('__') and hasattr(node, f):
+        new_fields[f] = self.copy(getattr(node, f))
     new_node = type(node)(**new_fields)
-    if anno.hasanno(node, anno.Basic.SKIP_PROCESSING):
-      anno.setanno(new_node, anno.Basic.SKIP_PROCESSING, True)
+
+    if self.preserve_annos:
+      for k in self.preserve_annos:
+        anno.copyanno(node, new_node, k)
     return new_node
 
 
-def copy_clean(node):
-  copier = CleanCopier()
-  if isinstance(node, list):
-    return [copier.visit(n) for n in node]
-  elif isinstance(node, tuple):
-    return tuple(copier.visit(n) for n in node)
-  else:
-    return copier.visit(node)
+def copy_clean(node, preserve_annos=None):
+  """Creates a deep copy of an AST.
+
+  The copy will not include fields that are prefixed by '__', with the
+  exception of user-specified annotations.
+
+  Args:
+    node: ast.AST
+    preserve_annos: Optional[Set[Hashable]], annotation keys to include in the
+        copy
+  Returns:
+    ast.AST
+  """
+  return CleanCopier(preserve_annos).copy(node)
 
 
 class SymbolRenamer(gast.NodeTransformer):
@@ -78,7 +84,11 @@ class SymbolRenamer(gast.NodeTransformer):
   def _process(self, node):
     qn = anno.getanno(node, anno.Basic.QN)
     if qn in self.name_map:
-      return gast.Name(str(self.name_map[qn]), node.ctx, None)
+      new_node = gast.Name(str(self.name_map[qn]), node.ctx, None)
+      # All annotations get carried over.
+      for k in anno.keys(node):
+        anno.copyanno(node, new_node, k)
+      return new_node
     return self.generic_visit(node)
 
   def visit_Name(self, node):
@@ -92,6 +102,7 @@ class SymbolRenamer(gast.NodeTransformer):
 
 
 def rename_symbols(node, name_map):
+  """Renames symbols in an AST. Requires qual_names annotations."""
   renamer = SymbolRenamer(name_map)
   if isinstance(node, list):
     return [renamer.visit(n) for n in node]
@@ -101,6 +112,7 @@ def rename_symbols(node, name_map):
 
 
 def keywords_to_dict(keywords):
+  """Converts a list of ast.keyword objects to a dict."""
   keys = []
   values = []
   for kw in keywords:
@@ -110,10 +122,7 @@ def keywords_to_dict(keywords):
 
 
 class PatternMatcher(gast.NodeVisitor):
-  """Matches a node against a pattern represented by a node.
-
-  The pattern may contain wildcards represented by the symbol '_'.
-  """
+  """Matches a node against a pattern represented by a node."""
 
   def __init__(self, pattern):
     self.pattern = pattern
@@ -177,9 +186,68 @@ class PatternMatcher(gast.NodeVisitor):
 
 
 def matches(node, pattern):
+  """Basic pattern matcher for AST.
+
+  The pattern may contain wildcards represented by the symbol '_'. A node
+  matches a pattern if for every node in the tree, either there is a node of
+  the same type in pattern, or a Name node with id='_'.
+
+  Args:
+    node: ast.AST
+    pattern: ast.AST
+  Returns:
+    bool
+  """
   if isinstance(pattern, str):
     pattern = parser.parse_expression(pattern)
   matcher = PatternMatcher(pattern)
   matcher.visit(node)
   return matcher.matches
 
+
+# TODO(mdan): Once we have error tracing, we may be able to just go to SSA.
+def apply_to_single_assignments(targets, values, apply_fn):
+  """Applies a function to each individual assignment.
+
+  This function can process a possibly-unpacked (e.g. a, b = c, d) assignment.
+  It tries to break down the unpacking if possible. In effect, it has the same
+  effect as passing the assigned values in SSA form to apply_fn.
+
+  Examples:
+
+  The following will result in apply_fn(a, c), apply_fn(b, d):
+
+      a, b = c, d
+
+  The following will result in apply_fn(a, c[0]), apply_fn(b, c[1]):
+
+      a, b = c
+
+  The following will result in apply_fn(a, (b, c)):
+
+      a = b, c
+
+  It uses the visitor pattern to allow subclasses to process single
+  assignments individually.
+
+  Args:
+    targets: Union[List[ast.AST, ...], Tuple[ast.AST, ...], ast.AST, should be
+        used with the targets field of an ast.Assign node
+    values: ast.AST
+    apply_fn: Callable[[ast.AST, ast.AST], None], called with the
+        respective nodes of each single assignment
+  """
+  if not isinstance(targets, (list, tuple)):
+    targets = (targets,)
+  for target in targets:
+    if isinstance(target, (gast.Tuple, gast.List)):
+      for i in range(len(target.elts)):
+        target_el = target.elts[i]
+        if isinstance(values, (gast.Tuple, gast.List)):
+          value_el = values.elts[i]
+        else:
+          idx = parser.parse_expression(str(i))
+          value_el = gast.Subscript(values, gast.Index(idx), ctx=gast.Load())
+        apply_to_single_assignments(target_el, value_el, apply_fn)
+    else:
+      apply_fn(target, values)
diff --git a/tensorflow/contrib/autograph/pyct/ast_util_test.py b/tensorflow/contrib/autograph/pyct/ast_util_test.py
index 3afa04a506..bd546c7f48 100644
--- a/tensorflow/contrib/autograph/pyct/ast_util_test.py
+++ b/tensorflow/contrib/autograph/pyct/ast_util_test.py
@@ -19,7 +19,10 @@ from __future__ import division
 from __future__ import print_function
 
 import ast
+import collections
+import textwrap
 
+from tensorflow.contrib.autograph.pyct import anno
 from tensorflow.contrib.autograph.pyct import ast_util
 from tensorflow.contrib.autograph.pyct import compiler
 from tensorflow.contrib.autograph.pyct import parser
@@ -29,53 +32,65 @@ from tensorflow.python.platform import test
 
 class AstUtilTest(test.TestCase):
 
-  def test_rename_symbols(self):
-    node = ast.Tuple([
-        ast.Name('a', ast.Load()),
-        ast.Name('b', ast.Load()),
-        ast.Attribute(ast.Name('b', None), 'c', ast.Store()),
-        ast.Attribute(
-            ast.Attribute(ast.Name('b', None), 'c', ast.Load()), 'd', None)
-    ], None)
+  def setUp(self):
+    super(AstUtilTest, self).setUp()
+    self._invocation_counts = collections.defaultdict(lambda: 0)
+
+  def test_rename_symbols_basic(self):
+    node = parser.parse_str('a + b')
+    node = qual_names.resolve(node)
+
+    node = ast_util.rename_symbols(
+        node, {qual_names.QN('a'): qual_names.QN('renamed_a')})
+
+    self.assertIsInstance(node.body[0].value.left.id, str)
+    self.assertEqual(compiler.ast_to_source(node).strip(), 'renamed_a + b')
+
+  def test_rename_symbols_attributes(self):
+    node = parser.parse_str('b.c = b.c.d')
     node = qual_names.resolve(node)
+
     node = ast_util.rename_symbols(
-        node, {
-            qual_names.QN('a'):
-                qual_names.QN('renamed_a'),
-            qual_names.QN(qual_names.QN('b'), attr='c'):
-                qual_names.QN('renamed_b_c'),
-        })
-
-    self.assertEqual(node.elts[0].id, 'renamed_a')
-    self.assertTrue(isinstance(node.elts[0].ctx, ast.Load))
-    self.assertEqual(node.elts[1].id, 'b')
-    self.assertEqual(node.elts[2].id, 'renamed_b_c')
-    self.assertTrue(isinstance(node.elts[2].ctx, ast.Store))
-    self.assertEqual(node.elts[3].value.id, 'renamed_b_c')
-    self.assertTrue(isinstance(node.elts[3].value.ctx, ast.Load))
+        node, {qual_names.from_str('b.c'): qual_names.QN('renamed_b_c')})
+
+    self.assertEqual(
+        compiler.ast_to_source(node).strip(), 'renamed_b_c = renamed_b_c.d')
+
+  def test_rename_symbols_annotations(self):
+    node = parser.parse_str('a[i]')
+    node = qual_names.resolve(node)
+    anno.setanno(node, 'foo', 'bar')
+    orig_anno = anno.getanno(node, 'foo')
+
+    node = ast_util.rename_symbols(node,
+                                   {qual_names.QN('a'): qual_names.QN('b')})
+
+    self.assertIs(anno.getanno(node, 'foo'), orig_anno)
 
   def test_copy_clean(self):
-    ret = ast.Return(
-        ast.BinOp(
-            op=ast.Add(),
-            left=ast.Name(id='a', ctx=ast.Load()),
-            right=ast.Num(1)))
-    setattr(ret, '__foo', 'bar')
-    node = ast.FunctionDef(
-        name='f',
-        args=ast.arguments(
-            args=[ast.Name(id='a', ctx=ast.Param())],
-            vararg=None,
-            kwarg=None,
-            defaults=[]),
-        body=[ret],
-        decorator_list=[],
-        returns=None)
+    node = parser.parse_str(
+        textwrap.dedent("""
+      def f(a):
+        return a + 1
+    """))
+    setattr(node.body[0], '__foo', 'bar')
     new_node = ast_util.copy_clean(node)
-    self.assertFalse(node is new_node)
-    self.assertFalse(ret is new_node.body[0])
+    self.assertIsNot(new_node, node)
+    self.assertIsNot(new_node.body[0], node.body[0])
     self.assertFalse(hasattr(new_node.body[0], '__foo'))
 
+  def test_copy_clean_preserves_annotations(self):
+    node = parser.parse_str(
+        textwrap.dedent("""
+      def f(a):
+        return a + 1
+    """))
+    anno.setanno(node.body[0], 'foo', 'bar')
+    anno.setanno(node.body[0], 'baz', 1)
+    new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
+    self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar')
+    self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
+
   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)
@@ -113,6 +128,33 @@ class AstUtilTest(test.TestCase):
     self.assertNoMatch('super(Foo, self).__init__()',
                        'super(Bar, _).__init__(_)')
 
+  def _mock_apply_fn(self, target, source):
+    target = compiler.ast_to_source(target).strip()
+    source = compiler.ast_to_source(source).strip()
+    self._invocation_counts[(target, source)] += 1
+
+  def test_apply_to_single_assignments_dynamic_unpack(self):
+    node = parser.parse_str('a, b, c = d')
+    node = node.body[0]
+    ast_util.apply_to_single_assignments(node.targets, node.value,
+                                         self._mock_apply_fn)
+    self.assertDictEqual(self._invocation_counts, {
+        ('a', 'd[0]'): 1,
+        ('b', 'd[1]'): 1,
+        ('c', 'd[2]'): 1,
+    })
+
+  def test_apply_to_single_assignments_static_unpack(self):
+    node = parser.parse_str('a, b, c = d, e, f')
+    node = node.body[0]
+    ast_util.apply_to_single_assignments(node.targets, node.value,
+                                         self._mock_apply_fn)
+    self.assertDictEqual(self._invocation_counts, {
+        ('a', 'd'): 1,
+        ('b', 'e'): 1,
+        ('c', 'f'): 1,
+    })
+
 
 if __name__ == '__main__':
   test.main()