diff options
Diffstat (limited to 'tensorflow/contrib/autograph/pyct/ast_util_test.py')
| -rw-r--r-- | tensorflow/contrib/autograph/pyct/ast_util_test.py | 142 |
1 files changed, 102 insertions, 40 deletions
diff --git a/tensorflow/contrib/autograph/pyct/ast_util_test.py b/tensorflow/contrib/autograph/pyct/ast_util_test.py index 3afa04a506..981e398b93 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,66 @@ 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) + source, _ = compiler.ast_to_source(node) + self.assertEqual(source.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')}) + + source, _ = compiler.ast_to_source(node) + self.assertEqual(source.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 +129,52 @@ 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) + source, _ = compiler.ast_to_source(source) + self._invocation_counts[(target.strip(), source.strip())] += 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, + }) + + def test_parallel_walk(self): + ret = ast.Return( + ast.BinOp( + op=ast.Add(), + left=ast.Name(id='a', ctx=ast.Load()), + right=ast.Num(1))) + 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) + for child_a, child_b in ast_util.parallel_walk(node, node): + self.assertEqual(child_a, child_b) + if __name__ == '__main__': test.main() |