xref: /aosp_15_r20/external/yapf/yapf/yapflib/identify_container.py (revision 7249d1a64f4850ccf838e62a46276f891f72998e) 
1# Copyright 2018 Google Inc. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Identify containers for lib2to3 trees.
15
16This module identifies containers and the elements in them. Each element points
17to the opening bracket and vice-versa.
18
19  IdentifyContainers(): the main function exported by this module.
20"""
21
22from lib2to3.pgen2 import token as grammar_token
23
24from yapf.yapflib import pytree_utils
25from yapf.yapflib import pytree_visitor
26
27
28def IdentifyContainers(tree):
29  """Run the identify containers visitor over the tree, modifying it in place.
30
31  Arguments:
32    tree: the top-level pytree node to annotate with subtypes.
33  """
34  identify_containers = _IdentifyContainers()
35  identify_containers.Visit(tree)
36
37
38class _IdentifyContainers(pytree_visitor.PyTreeVisitor):
39  """_IdentifyContainers - see file-level docstring for detailed description."""
40
41  def Visit_trailer(self, node):  # pylint: disable=invalid-name
42    for child in node.children:
43      self.Visit(child)
44
45    if len(node.children) != 3:
46      return
47    if node.children[0].type != grammar_token.LPAR:
48      return
49
50    if pytree_utils.NodeName(node.children[1]) == 'arglist':
51      for child in node.children[1].children:
52        pytree_utils.SetOpeningBracket(
53            pytree_utils.FirstLeafNode(child), node.children[0])
54    else:
55      pytree_utils.SetOpeningBracket(
56          pytree_utils.FirstLeafNode(node.children[1]), node.children[0])
57
58  def Visit_atom(self, node):  # pylint: disable=invalid-name
59    for child in node.children:
60      self.Visit(child)
61
62    if len(node.children) != 3:
63      return
64    if node.children[0].type != grammar_token.LPAR:
65      return
66
67    for child in node.children[1].children:
68      pytree_utils.SetOpeningBracket(
69          pytree_utils.FirstLeafNode(child), node.children[0])
70