# Copyright 2015 Google Inc. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for yapf.pytree_utils.""" import unittest from lib2to3 import pygram from lib2to3 import pytree from lib2to3.pgen2 import token from yapf.yapflib import pytree_utils # More direct access to the symbol->number mapping living within the grammar # module. _GRAMMAR_SYMBOL2NUMBER = pygram.python_grammar.symbol2number _FOO = 'foo' _FOO1 = 'foo1' _FOO2 = 'foo2' _FOO3 = 'foo3' _FOO4 = 'foo4' _FOO5 = 'foo5' class NodeNameTest(unittest.TestCase): def testNodeNameForLeaf(self): leaf = pytree.Leaf(token.LPAR, '(') self.assertEqual('LPAR', pytree_utils.NodeName(leaf)) def testNodeNameForNode(self): leaf = pytree.Leaf(token.LPAR, '(') node = pytree.Node(pygram.python_grammar.symbol2number['suite'], [leaf]) self.assertEqual('suite', pytree_utils.NodeName(node)) class ParseCodeToTreeTest(unittest.TestCase): def testParseCodeToTree(self): # Since ParseCodeToTree is a thin wrapper around underlying lib2to3 # functionality, only a sanity test here... tree = pytree_utils.ParseCodeToTree('foo = 2\n') self.assertEqual('file_input', pytree_utils.NodeName(tree)) self.assertEqual(2, len(tree.children)) self.assertEqual('simple_stmt', pytree_utils.NodeName(tree.children[0])) def testPrintFunctionToTree(self): tree = pytree_utils.ParseCodeToTree( 'print("hello world", file=sys.stderr)\n') self.assertEqual('file_input', pytree_utils.NodeName(tree)) self.assertEqual(2, len(tree.children)) self.assertEqual('simple_stmt', pytree_utils.NodeName(tree.children[0])) def testPrintStatementToTree(self): tree = pytree_utils.ParseCodeToTree('print "hello world"\n') self.assertEqual('file_input', pytree_utils.NodeName(tree)) self.assertEqual(2, len(tree.children)) self.assertEqual('simple_stmt', pytree_utils.NodeName(tree.children[0])) def testClassNotLocal(self): tree = pytree_utils.ParseCodeToTree('class nonlocal: pass\n') self.assertEqual('file_input', pytree_utils.NodeName(tree)) self.assertEqual(2, len(tree.children)) self.assertEqual('classdef', pytree_utils.NodeName(tree.children[0])) class InsertNodesBeforeAfterTest(unittest.TestCase): def _BuildSimpleTree(self): # Builds a simple tree we can play with in the tests. # The tree looks like this: # # suite: # LPAR # LPAR # simple_stmt: # NAME('foo') # lpar1 = pytree.Leaf(token.LPAR, '(') lpar2 = pytree.Leaf(token.LPAR, '(') simple_stmt = pytree.Node(_GRAMMAR_SYMBOL2NUMBER['simple_stmt'], [pytree.Leaf(token.NAME, 'foo')]) return pytree.Node(_GRAMMAR_SYMBOL2NUMBER['suite'], [lpar1, lpar2, simple_stmt]) def _MakeNewNodeRPAR(self): return pytree.Leaf(token.RPAR, ')') def setUp(self): self._simple_tree = self._BuildSimpleTree() def testInsertNodesBefore(self): # Insert before simple_stmt and make sure it went to the right place pytree_utils.InsertNodesBefore([self._MakeNewNodeRPAR()], self._simple_tree.children[2]) self.assertEqual(4, len(self._simple_tree.children)) self.assertEqual('RPAR', pytree_utils.NodeName(self._simple_tree.children[2])) self.assertEqual('simple_stmt', pytree_utils.NodeName(self._simple_tree.children[3])) def testInsertNodesBeforeFirstChild(self): # Insert before the first child of its parent simple_stmt = self._simple_tree.children[2] foo_child = simple_stmt.children[0] pytree_utils.InsertNodesBefore([self._MakeNewNodeRPAR()], foo_child) self.assertEqual(3, len(self._simple_tree.children)) self.assertEqual(2, len(simple_stmt.children)) self.assertEqual('RPAR', pytree_utils.NodeName(simple_stmt.children[0])) self.assertEqual('NAME', pytree_utils.NodeName(simple_stmt.children[1])) def testInsertNodesAfter(self): # Insert after and make sure it went to the right place pytree_utils.InsertNodesAfter([self._MakeNewNodeRPAR()], self._simple_tree.children[2]) self.assertEqual(4, len(self._simple_tree.children)) self.assertEqual('simple_stmt', pytree_utils.NodeName(self._simple_tree.children[2])) self.assertEqual('RPAR', pytree_utils.NodeName(self._simple_tree.children[3])) def testInsertNodesAfterLastChild(self): # Insert after the last child of its parent simple_stmt = self._simple_tree.children[2] foo_child = simple_stmt.children[0] pytree_utils.InsertNodesAfter([self._MakeNewNodeRPAR()], foo_child) self.assertEqual(3, len(self._simple_tree.children)) self.assertEqual(2, len(simple_stmt.children)) self.assertEqual('NAME', pytree_utils.NodeName(simple_stmt.children[0])) self.assertEqual('RPAR', pytree_utils.NodeName(simple_stmt.children[1])) def testInsertNodesWhichHasParent(self): # Try to insert an existing tree node into another place and fail. with self.assertRaises(RuntimeError): pytree_utils.InsertNodesAfter([self._simple_tree.children[1]], self._simple_tree.children[0]) class AnnotationsTest(unittest.TestCase): def setUp(self): self._leaf = pytree.Leaf(token.LPAR, '(') self._node = pytree.Node(_GRAMMAR_SYMBOL2NUMBER['simple_stmt'], [pytree.Leaf(token.NAME, 'foo')]) def testGetWhenNone(self): self.assertIsNone(pytree_utils.GetNodeAnnotation(self._leaf, _FOO)) def testSetWhenNone(self): pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO), 20) def testSetAgain(self): pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO), 20) pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 30) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO), 30) def testMultiple(self): pytree_utils.SetNodeAnnotation(self._leaf, _FOO, 20) pytree_utils.SetNodeAnnotation(self._leaf, _FOO1, 1) pytree_utils.SetNodeAnnotation(self._leaf, _FOO2, 2) pytree_utils.SetNodeAnnotation(self._leaf, _FOO3, 3) pytree_utils.SetNodeAnnotation(self._leaf, _FOO4, 4) pytree_utils.SetNodeAnnotation(self._leaf, _FOO5, 5) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO), 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO1), 1) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO2), 2) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO3), 3) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO4), 4) self.assertEqual(pytree_utils.GetNodeAnnotation(self._leaf, _FOO5), 5) def testSubtype(self): pytree_utils.AppendNodeAnnotation(self._leaf, pytree_utils.Annotation.SUBTYPE, _FOO) self.assertSetEqual( pytree_utils.GetNodeAnnotation(self._leaf, pytree_utils.Annotation.SUBTYPE), {_FOO}) pytree_utils.RemoveSubtypeAnnotation(self._leaf, _FOO) self.assertSetEqual( pytree_utils.GetNodeAnnotation(self._leaf, pytree_utils.Annotation.SUBTYPE), set()) def testSetOnNode(self): pytree_utils.SetNodeAnnotation(self._node, _FOO, 20) self.assertEqual(pytree_utils.GetNodeAnnotation(self._node, _FOO), 20) if __name__ == '__main__': unittest.main()