aboutsummaryrefslogtreecommitdiff
path: root/yapf/yapflib/pytree_unwrapper.py
diff options
context:
space:
mode:
Diffstat (limited to 'yapf/yapflib/pytree_unwrapper.py')
-rw-r--r--yapf/yapflib/pytree_unwrapper.py117
1 files changed, 77 insertions, 40 deletions
diff --git a/yapf/yapflib/pytree_unwrapper.py b/yapf/yapflib/pytree_unwrapper.py
index 0d371ae..1b05b0e 100644
--- a/yapf/yapflib/pytree_unwrapper.py
+++ b/yapf/yapflib/pytree_unwrapper.py
@@ -11,13 +11,13 @@
# 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.
-"""PyTreeUnwrapper - produces a list of unwrapped lines from a pytree.
+"""PyTreeUnwrapper - produces a list of logical lines from a pytree.
-[for a description of what an unwrapped line is, see unwrapped_line.py]
+[for a description of what a logical line is, see logical_line.py]
This is a pytree visitor that goes over a parse tree and produces a list of
-UnwrappedLine containers from it, each with its own depth and containing all
-the tokens that could fit on the line if there were no maximal line-length
+LogicalLine containers from it, each with its own depth and containing all the
+tokens that could fit on the line if there were no maximal line-length
limitations.
Note: a precondition to running this visitor and obtaining correct results is
@@ -31,27 +31,30 @@ For most uses, the convenience function UnwrapPyTree should be sufficient.
from lib2to3 import pytree
from lib2to3.pgen2 import token as grammar_token
+from yapf.yapflib import format_token
+from yapf.yapflib import logical_line
+from yapf.yapflib import object_state
from yapf.yapflib import pytree_utils
from yapf.yapflib import pytree_visitor
from yapf.yapflib import split_penalty
from yapf.yapflib import style
-from yapf.yapflib import unwrapped_line
+from yapf.yapflib import subtypes
def UnwrapPyTree(tree):
- """Create and return a list of unwrapped lines from the given pytree.
+ """Create and return a list of logical lines from the given pytree.
Arguments:
- tree: the top-level pytree node to unwrap.
+ tree: the top-level pytree node to unwrap..
Returns:
- A list of UnwrappedLine objects.
+ A list of LogicalLine objects.
"""
unwrapper = PyTreeUnwrapper()
unwrapper.Visit(tree)
- uwlines = unwrapper.GetUnwrappedLines()
- uwlines.sort(key=lambda x: x.lineno)
- return uwlines
+ llines = unwrapper.GetLogicalLines()
+ llines.sort(key=lambda x: x.lineno)
+ return llines
# Grammar tokens considered as whitespace for the purpose of unwrapping.
@@ -77,39 +80,40 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
"""
def __init__(self):
- # A list of all unwrapped lines finished visiting so far.
- self._unwrapped_lines = []
+ # A list of all logical lines finished visiting so far.
+ self._logical_lines = []
- # Builds up a "current" unwrapped line while visiting pytree nodes. Some
- # nodes will finish a line and start a new one.
- self._cur_unwrapped_line = unwrapped_line.UnwrappedLine(0)
+ # Builds up a "current" logical line while visiting pytree nodes. Some nodes
+ # will finish a line and start a new one.
+ self._cur_logical_line = logical_line.LogicalLine(0)
# Current indentation depth.
self._cur_depth = 0
- def GetUnwrappedLines(self):
+ def GetLogicalLines(self):
"""Fetch the result of the tree walk.
Note: only call this after visiting the whole tree.
Returns:
- A list of UnwrappedLine objects.
+ A list of LogicalLine objects.
"""
# Make sure the last line that was being populated is flushed.
self._StartNewLine()
- return self._unwrapped_lines
+ return self._logical_lines
def _StartNewLine(self):
"""Finish current line and start a new one.
- Place the currently accumulated line into the _unwrapped_lines list and
+ Place the currently accumulated line into the _logical_lines list and
start a new one.
"""
- if self._cur_unwrapped_line.tokens:
- self._unwrapped_lines.append(self._cur_unwrapped_line)
- _MatchBrackets(self._cur_unwrapped_line)
- _AdjustSplitPenalty(self._cur_unwrapped_line)
- self._cur_unwrapped_line = unwrapped_line.UnwrappedLine(self._cur_depth)
+ if self._cur_logical_line.tokens:
+ self._logical_lines.append(self._cur_logical_line)
+ _MatchBrackets(self._cur_logical_line)
+ _IdentifyParameterLists(self._cur_logical_line)
+ _AdjustSplitPenalty(self._cur_logical_line)
+ self._cur_logical_line = logical_line.LogicalLine(self._cur_depth)
_STMT_TYPES = frozenset({
'if_stmt',
@@ -148,7 +152,7 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
"""Helper for visiting compound statements.
Python compound statements serve as containers for other statements. Thus,
- when we encounter a new compound statement we start a new unwrapped line.
+ when we encounter a new compound statement, we start a new logical line.
Arguments:
node: the node to visit.
@@ -202,7 +206,7 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
for child in node.children:
index += 1
self.Visit(child)
- if pytree_utils.NodeName(child) == 'ASYNC':
+ if child.type == grammar_token.ASYNC:
break
for child in node.children[index].children:
self.Visit(child)
@@ -218,16 +222,17 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
for child in node.children:
index += 1
self.Visit(child)
- if pytree_utils.NodeName(child) == 'ASYNC':
+ if child.type == grammar_token.ASYNC:
break
for child in node.children[index].children:
+ if child.type == grammar_token.NAME and child.value == 'else':
+ self._StartNewLine()
self.Visit(child)
def Visit_decorator(self, node): # pylint: disable=invalid-name
for child in node.children:
self.Visit(child)
- if (pytree_utils.NodeName(child) == 'COMMENT' and
- child == node.children[0]):
+ if child.type == grammar_token.COMMENT and child == node.children[0]:
self._StartNewLine()
def Visit_decorators(self, node): # pylint: disable=invalid-name
@@ -280,7 +285,7 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
def DefaultLeafVisit(self, leaf):
"""Default visitor for tree leaves.
- A tree leaf is always just gets appended to the current unwrapped line.
+ A tree leaf is always just gets appended to the current logical line.
Arguments:
leaf: the leaf to visit.
@@ -289,13 +294,13 @@ class PyTreeUnwrapper(pytree_visitor.PyTreeVisitor):
self._StartNewLine()
elif leaf.type != grammar_token.COMMENT or leaf.value.strip():
# Add non-whitespace tokens and comments that aren't empty.
- self._cur_unwrapped_line.AppendNode(leaf)
+ self._cur_logical_line.AppendNode(leaf)
_BRACKET_MATCH = {')': '(', '}': '{', ']': '['}
-def _MatchBrackets(uwline):
+def _MatchBrackets(line):
"""Visit the node and match the brackets.
For every open bracket ('[', '{', or '('), find the associated closing bracket
@@ -303,10 +308,10 @@ def _MatchBrackets(uwline):
or close bracket.
Arguments:
- uwline: (UnwrappedLine) An unwrapped line.
+ line: (LogicalLine) A logical line.
"""
bracket_stack = []
- for token in uwline.tokens:
+ for token in line.tokens:
if token.value in pytree_utils.OPENING_BRACKETS:
bracket_stack.append(token)
elif token.value in pytree_utils.CLOSING_BRACKETS:
@@ -320,17 +325,50 @@ def _MatchBrackets(uwline):
token.container_opening = bracket
-def _AdjustSplitPenalty(uwline):
+def _IdentifyParameterLists(line):
+ """Visit the node to create a state for parameter lists.
+
+ For instance, a parameter is considered an "object" with its first and last
+ token uniquely identifying the object.
+
+ Arguments:
+ line: (LogicalLine) A logical line.
+ """
+ func_stack = []
+ param_stack = []
+ for tok in line.tokens:
+ # Identify parameter list objects.
+ if subtypes.FUNC_DEF in tok.subtypes:
+ assert tok.next_token.value == '('
+ func_stack.append(tok.next_token)
+ continue
+
+ if func_stack and tok.value == ')':
+ if tok == func_stack[-1].matching_bracket:
+ func_stack.pop()
+ continue
+
+ # Identify parameter objects.
+ if subtypes.PARAMETER_START in tok.subtypes:
+ param_stack.append(tok)
+
+ # Not "elif", a parameter could be a single token.
+ if param_stack and subtypes.PARAMETER_STOP in tok.subtypes:
+ start = param_stack.pop()
+ func_stack[-1].parameters.append(object_state.Parameter(start, tok))
+
+
+def _AdjustSplitPenalty(line):
"""Visit the node and adjust the split penalties if needed.
A token shouldn't be split if it's not within a bracket pair. Mark any token
that's not within a bracket pair as "unbreakable".
Arguments:
- uwline: (UnwrappedLine) An unwrapped line.
+ line: (LogicalLine) An logical line.
"""
bracket_level = 0
- for index, token in enumerate(uwline.tokens):
+ for index, token in enumerate(line.tokens):
if index and not bracket_level:
pytree_utils.SetNodeAnnotation(token.node,
pytree_utils.Annotation.SPLIT_PENALTY,
@@ -348,8 +386,7 @@ def _DetermineMustSplitAnnotation(node):
if not _ContainsComments(node):
token = next(node.parent.leaves())
if token.value == '(':
- if sum(1 for ch in node.children
- if pytree_utils.NodeName(ch) == 'COMMA') < 2:
+ if sum(1 for ch in node.children if ch.type == grammar_token.COMMA) < 2:
return
if (not isinstance(node.children[-1], pytree.Leaf) or
node.children[-1].value != ','):