xref: /aosp_15_r20/external/yapf/yapf/yapflib/blank_line_calculator.py (revision 7249d1a64f4850ccf838e62a46276f891f72998e) 
1# Copyright 2015 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"""Calculate the number of blank lines between top-level entities.
15
16Calculates how many blank lines we need between classes, functions, and other
17entities at the same level.
18
19  CalculateBlankLines(): the main function exported by this module.
20
21Annotations:
22  newlines: The number of newlines required before the node.
23"""
24
25from lib2to3.pgen2 import token as grammar_token
26
27from yapf.yapflib import py3compat
28from yapf.yapflib import pytree_utils
29from yapf.yapflib import pytree_visitor
30from yapf.yapflib import style
31
32_NO_BLANK_LINES = 1
33_ONE_BLANK_LINE = 2
34_TWO_BLANK_LINES = 3
35
36_PYTHON_STATEMENTS = frozenset({
37    'small_stmt', 'expr_stmt', 'print_stmt', 'del_stmt', 'pass_stmt',
38    'break_stmt', 'continue_stmt', 'return_stmt', 'raise_stmt', 'yield_stmt',
39    'import_stmt', 'global_stmt', 'exec_stmt', 'assert_stmt', 'if_stmt',
40    'while_stmt', 'for_stmt', 'try_stmt', 'with_stmt', 'nonlocal_stmt',
41    'async_stmt', 'simple_stmt'
42})
43
44
45def CalculateBlankLines(tree):
46  """Run the blank line calculator visitor over the tree.
47
48  This modifies the tree in place.
49
50  Arguments:
51    tree: the top-level pytree node to annotate with subtypes.
52  """
53  blank_line_calculator = _BlankLineCalculator()
54  blank_line_calculator.Visit(tree)
55
56
57class _BlankLineCalculator(pytree_visitor.PyTreeVisitor):
58  """_BlankLineCalculator - see file-level docstring for a description."""
59
60  def __init__(self):
61    self.class_level = 0
62    self.function_level = 0
63    self.last_comment_lineno = 0
64    self.last_was_decorator = False
65    self.last_was_class_or_function = False
66
67  def Visit_simple_stmt(self, node):  # pylint: disable=invalid-name
68    self.DefaultNodeVisit(node)
69    if node.children[0].type == grammar_token.COMMENT:
70      self.last_comment_lineno = node.children[0].lineno
71
72  def Visit_decorator(self, node):  # pylint: disable=invalid-name
73    if (self.last_comment_lineno and
74        self.last_comment_lineno == node.children[0].lineno - 1):
75      _SetNumNewlines(node.children[0], _NO_BLANK_LINES)
76    else:
77      _SetNumNewlines(node.children[0], self._GetNumNewlines(node))
78    for child in node.children:
79      self.Visit(child)
80    self.last_was_decorator = True
81
82  def Visit_classdef(self, node):  # pylint: disable=invalid-name
83    self.last_was_class_or_function = False
84    index = self._SetBlankLinesBetweenCommentAndClassFunc(node)
85    self.last_was_decorator = False
86    self.class_level += 1
87    for child in node.children[index:]:
88      self.Visit(child)
89    self.class_level -= 1
90    self.last_was_class_or_function = True
91
92  def Visit_funcdef(self, node):  # pylint: disable=invalid-name
93    self.last_was_class_or_function = False
94    index = self._SetBlankLinesBetweenCommentAndClassFunc(node)
95    if _AsyncFunction(node):
96      index = self._SetBlankLinesBetweenCommentAndClassFunc(
97          node.prev_sibling.parent)
98      _SetNumNewlines(node.children[0], None)
99    else:
100      index = self._SetBlankLinesBetweenCommentAndClassFunc(node)
101    self.last_was_decorator = False
102    self.function_level += 1
103    for child in node.children[index:]:
104      self.Visit(child)
105    self.function_level -= 1
106    self.last_was_class_or_function = True
107
108  def DefaultNodeVisit(self, node):
109    """Override the default visitor for Node.
110
111    This will set the blank lines required if the last entity was a class or
112    function.
113
114    Arguments:
115      node: (pytree.Node) The node to visit.
116    """
117    if self.last_was_class_or_function:
118      if pytree_utils.NodeName(node) in _PYTHON_STATEMENTS:
119        leaf = pytree_utils.FirstLeafNode(node)
120        _SetNumNewlines(leaf, self._GetNumNewlines(leaf))
121    self.last_was_class_or_function = False
122    super(_BlankLineCalculator, self).DefaultNodeVisit(node)
123
124  def _SetBlankLinesBetweenCommentAndClassFunc(self, node):
125    """Set the number of blanks between a comment and class or func definition.
126
127    Class and function definitions have leading comments as children of the
128    classdef and functdef nodes.
129
130    Arguments:
131      node: (pytree.Node) The classdef or funcdef node.
132
133    Returns:
134      The index of the first child past the comment nodes.
135    """
136    index = 0
137    while pytree_utils.IsCommentStatement(node.children[index]):
138      # Standalone comments are wrapped in a simple_stmt node with the comment
139      # node as its only child.
140      self.Visit(node.children[index].children[0])
141      if not self.last_was_decorator:
142        _SetNumNewlines(node.children[index].children[0], _ONE_BLANK_LINE)
143      index += 1
144    if (index and node.children[index].lineno - 1
145        == node.children[index - 1].children[0].lineno):
146      _SetNumNewlines(node.children[index], _NO_BLANK_LINES)
147    else:
148      if self.last_comment_lineno + 1 == node.children[index].lineno:
149        num_newlines = _NO_BLANK_LINES
150      else:
151        num_newlines = self._GetNumNewlines(node)
152      _SetNumNewlines(node.children[index], num_newlines)
153    return index
154
155  def _GetNumNewlines(self, node):
156    if self.last_was_decorator:
157      return _NO_BLANK_LINES
158    elif self._IsTopLevel(node):
159      return 1 + style.Get('BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION')
160    return _ONE_BLANK_LINE
161
162  def _IsTopLevel(self, node):
163    return (not (self.class_level or self.function_level) and
164            _StartsInZerothColumn(node))
165
166
167def _SetNumNewlines(node, num_newlines):
168  pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.NEWLINES,
169                                 num_newlines)
170
171
172def _StartsInZerothColumn(node):
173  return (pytree_utils.FirstLeafNode(node).column == 0 or
174          (_AsyncFunction(node) and node.prev_sibling.column == 0))
175
176
177def _AsyncFunction(node):
178  return (py3compat.PY3 and node.prev_sibling and
179          node.prev_sibling.type == grammar_token.ASYNC)
180