xref: /aosp_15_r20/external/yapf/yapftests/pytree_utils_test.py (revision 7249d1a64f4850ccf838e62a46276f891f72998e)

# 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()