• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 yapf.yapflib import pytree_utils
23from yapf.yapflib import pytree_visitor
24
25
26def IdentifyContainers(tree):
27  """Run the identify containers visitor over the tree, modifying it in place.
28
29  Arguments:
30    tree: the top-level pytree node to annotate with subtypes.
31  """
32  identify_containers = _IdentifyContainers()
33  identify_containers.Visit(tree)
34
35
36class _IdentifyContainers(pytree_visitor.PyTreeVisitor):
37  """_IdentifyContainers - see file-level docstring for detailed description."""
38
39  def Visit_trailer(self, node):  # pylint: disable=invalid-name
40    for child in node.children:
41      self.Visit(child)
42
43    if len(node.children) != 3:
44      return
45    if pytree_utils.NodeName(node.children[0]) != 'LPAR':
46      return
47
48    if pytree_utils.NodeName(node.children[1]) == 'arglist':
49      for child in node.children[1].children:
50        pytree_utils.SetOpeningBracket(
51            pytree_utils.FirstLeafNode(child), node.children[0])
52    else:
53      pytree_utils.SetOpeningBracket(
54          pytree_utils.FirstLeafNode(node.children[1]), node.children[0])
55
56  def Visit_atom(self, node):  # pylint: disable=invalid-name
57    for child in node.children:
58      self.Visit(child)
59
60    if len(node.children) != 3:
61      return
62    if pytree_utils.NodeName(node.children[0]) != 'LPAR':
63      return
64
65    for child in node.children[1].children:
66      pytree_utils.SetOpeningBracket(
67          pytree_utils.FirstLeafNode(child), node.children[0])
68