diff --git a/python/src/com/jetbrains/python/formatter/PyBinaryExpressionBlock.java b/python/src/com/jetbrains/python/formatter/PyBinaryExpressionBlock.java new file mode 100644 index 000000000000..ed6a9c5a10e0 --- /dev/null +++ b/python/src/com/jetbrains/python/formatter/PyBinaryExpressionBlock.java @@ -0,0 +1,57 @@ +// Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +package com.jetbrains.python.formatter; + +import com.intellij.formatting.Alignment; +import com.intellij.formatting.Indent; +import com.intellij.formatting.Wrap; +import com.intellij.lang.ASTNode; +import com.jetbrains.python.PyElementTypes; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +import java.util.ArrayList; +import java.util.List; + +/** + * @author Mikhail Golubev + */ +public class PyBinaryExpressionBlock extends PyBlock { + + private List myChildrenNodes = new ArrayList<>(); + + public PyBinaryExpressionBlock(@Nullable PyBlock parent, + @NotNull ASTNode node, + @Nullable Alignment alignment, + @NotNull Indent indent, + @Nullable Wrap wrap, + @NotNull PyBlockContext context) { + super(parent, node, alignment, indent, wrap, context); + assert node.getElementType() == PyElementTypes.BINARY_EXPRESSION; + + myChildrenNodes = collectChildren(); + } + + @NotNull + private List collectChildren() { + List result = new ArrayList<>(); + collectChildren(myNode, result); + return result; + } + + private void collectChildren(@NotNull ASTNode node, @NotNull List result) { + if (node.getElementType() == PyElementTypes.BINARY_EXPRESSION) { + for (ASTNode child : node.getChildren(null)) { + collectChildren(child, result); + } + } + else if (node != this) { + result.add(node); + } + } + + @NotNull + @Override + protected Iterable getSubBlockNodes() { + return myChildrenNodes; + } +} diff --git a/python/src/com/jetbrains/python/formatter/PyBlock.java b/python/src/com/jetbrains/python/formatter/PyBlock.java index ad926710e75d..c275cb567223 100644 --- a/python/src/com/jetbrains/python/formatter/PyBlock.java +++ b/python/src/com/jetbrains/python/formatter/PyBlock.java @@ -81,10 +81,23 @@ public class PyBlock implements ASTBlock { public static final Key IMPORT_GROUP_BEGIN = Key.create("com.jetbrains.python.formatter.importGroupBegin"); private static final boolean ALIGN_IF_CONDITION_WITHOUT_PARENTHESES = false; + @NotNull + public static PyBlock createBlock(@Nullable PyBlock parent, + @NotNull ASTNode node, + @Nullable Alignment alignment, + @NotNull Indent indent, + @Nullable Wrap wrap, + @NotNull PyBlockContext context) { + if (node.getElementType() == PyElementTypes.BINARY_EXPRESSION) { + return new PyBinaryExpressionBlock(parent, node, alignment, indent, wrap, context); + } + return new PyBlock(parent, node, alignment, indent, wrap, context); + } + private final PyBlock myParent; private final Alignment myAlignment; private final Indent myIndent; - private final ASTNode myNode; + protected final ASTNode myNode; private final Wrap myWrap; private final PyBlockContext myContext; private List mySubBlocks = null; @@ -97,12 +110,12 @@ public class PyBlock implements ASTBlock { private Wrap myDictWrapping = null; private Wrap myFromImportWrapping = null; - public PyBlock(@Nullable PyBlock parent, - @NotNull ASTNode node, - @Nullable Alignment alignment, - @NotNull Indent indent, - @Nullable Wrap wrap, - @NotNull PyBlockContext context) { + protected PyBlock(@Nullable PyBlock parent, + @NotNull ASTNode node, + @Nullable Alignment alignment, + @NotNull Indent indent, + @Nullable Wrap wrap, + @NotNull PyBlockContext context) { myParent = parent; myAlignment = alignment; myIndent = indent; @@ -163,7 +176,7 @@ public class PyBlock implements ASTBlock { @NotNull private Map buildSubBlocks() { final Map blocks = new LinkedHashMap<>(); - for (ASTNode child = myNode.getFirstChildNode(); child != null; child = child.getTreeNext()) { + for (ASTNode child: getSubBlockNodes()) { final IElementType childType = child.getElementType(); @@ -178,6 +191,11 @@ public class PyBlock implements ASTBlock { return Collections.unmodifiableMap(blocks); } + @NotNull + protected Iterable getSubBlockNodes() { + return Arrays.asList(myNode.getChildren(null)); + } + @NotNull private PyBlock buildSubBlock(@NotNull ASTNode child) { final IElementType parentType = myNode.getElementType(); @@ -405,7 +423,7 @@ public class PyBlock implements ASTBlock { prev = prev.getTreePrev(); } - return new PyBlock(this, child, childAlignment, childIndent, childWrap, myContext); + return createBlock(this, child, childAlignment, childIndent, childWrap, myContext); } private static boolean isParenthesisedIfCondition(@NotNull ASTNode node) { diff --git a/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java b/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java index 14851b412397..8cc865efc3a0 100644 --- a/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java +++ b/python/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java @@ -51,7 +51,7 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilderEx, C printAST(fileNode, 0); } final PyBlockContext context = new PyBlockContext(settings, createSpacingBuilder(settings), mode); - final PyBlock block = new PyBlock(null, element.getNode(), null, Indent.getNoneIndent(), null, context); + final PyBlock block = PyBlock.createBlock(null, element.getNode(), null, Indent.getNoneIndent(), null, context); if (DUMP_FORMATTING_AST) { FormattingModelDumper.dumpFormattingModel(block, 2, System.out); } diff --git a/python/testData/formatter/multilineElifCondition_after.py b/python/testData/formatter/multilineElifCondition_after.py index fe7b29f94299..d159919229c7 100644 --- a/python/testData/formatter/multilineElifCondition_after.py +++ b/python/testData/formatter/multilineElifCondition_after.py @@ -1,5 +1,5 @@ if foo and \ - bar == 42: + bar == 42: pass elif foo and \ bar: