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