You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
206 lines
7.9 KiB
206 lines
7.9 KiB
# 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()
|