python/src/com/jetbrains/python/formatter/PyBlock.java

   1 /*
   2  * Copyright 2000-2014 JetBrains s.r.o.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  * http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 package com.jetbrains.python.formatter;
  17
  18 import com.intellij.formatting.*;
  19 import com.intellij.lang.ASTNode;
  20 import com.intellij.openapi.editor.Document;
  21 import com.intellij.openapi.util.Key;
  22 import com.intellij.openapi.util.TextRange;
  23 import com.intellij.psi.*;
  24 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
  25 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
  26 import com.intellij.psi.impl.source.tree.TreeUtil;
  27 import com.intellij.psi.tree.IElementType;
  28 import com.intellij.psi.tree.TokenSet;
  29 import com.intellij.psi.util.PsiTreeUtil;
  30 import com.intellij.util.ArrayUtil;
  31 import com.jetbrains.python.PyElementTypes;
  32 import com.jetbrains.python.PyTokenTypes;
  33 import com.jetbrains.python.PythonDialectsTokenSetProvider;
  34 import com.jetbrains.python.psi.*;
  35 import com.jetbrains.python.psi.impl.PyPsiUtils;
  36 import org.jetbrains.annotations.NotNull;
  37 import org.jetbrains.annotations.Nullable;
  38
  39 import java.util.*;
  40
  41 import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_COLON;
  42 import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_VALUE;
  43 import static com.jetbrains.python.formatter.PythonFormattingModelBuilder.STATEMENT_OR_DECLARATION;
  44 import static com.jetbrains.python.psi.PyUtil.as;
  45
  46 /**
  47  * @author yole
  48  */
  49 public class PyBlock implements ASTBlock {
  50   private static final TokenSet ourListElementTypes = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION,
  51                                                                       PyElementTypes.LIST_COMP_EXPRESSION,
  52                                                                       PyElementTypes.DICT_COMP_EXPRESSION,
  53                                                                       PyElementTypes.SET_COMP_EXPRESSION,
  54                                                                       PyElementTypes.DICT_LITERAL_EXPRESSION,
  55                                                                       PyElementTypes.SET_LITERAL_EXPRESSION,
  56                                                                       PyElementTypes.ARGUMENT_LIST,
  57                                                                       PyElementTypes.PARAMETER_LIST,
  58                                                                       PyElementTypes.TUPLE_EXPRESSION,
  59                                                                       PyElementTypes.PARENTHESIZED_EXPRESSION,
  60                                                                       PyElementTypes.SLICE_EXPRESSION,
  61                                                                       PyElementTypes.SUBSCRIPTION_EXPRESSION,
  62                                                                       PyElementTypes.GENERATOR_EXPRESSION);
  63
  64   private static final TokenSet ourBrackets = TokenSet.create(PyTokenTypes.LPAR, PyTokenTypes.RPAR,
  65                                                               PyTokenTypes.LBRACE, PyTokenTypes.RBRACE,
  66                                                               PyTokenTypes.LBRACKET, PyTokenTypes.RBRACKET);
  67
  68   private static final TokenSet ourHangingIndentOwners = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION,
  69                                                                          PyElementTypes.DICT_LITERAL_EXPRESSION,
  70                                                                          PyElementTypes.SET_LITERAL_EXPRESSION,
  71                                                                          PyElementTypes.ARGUMENT_LIST,
  72                                                                          PyElementTypes.PARAMETER_LIST,
  73                                                                          PyElementTypes.TUPLE_EXPRESSION,
  74                                                                          PyElementTypes.PARENTHESIZED_EXPRESSION,
  75                                                                          PyElementTypes.GENERATOR_EXPRESSION,
  76                                                                          PyElementTypes.FUNCTION_DECLARATION,
  77                                                                          PyElementTypes.CALL_EXPRESSION,
  78                                                                          PyElementTypes.FROM_IMPORT_STATEMENT);
  79
  80   public static final Key<Boolean> IMPORT_GROUP_BEGIN = Key.create("com.jetbrains.python.formatter.importGroupBegin");
  81
  82   private final PyBlock myParent;
  83   private final Alignment myAlignment;
  84   private final Indent myIndent;
  85   private final ASTNode myNode;
  86   private final Wrap myWrap;
  87   private final PyBlockContext myContext;
  88   private List<PyBlock> mySubBlocks = null;
  89   private Map<ASTNode, PyBlock> mySubBlockByNode = null;
  90   private Alignment myChildAlignment;
  91   private final Alignment myDictAlignment;
  92   private final Wrap myDictWrapping;
  93   private final boolean myEmptySequence;
  94
  95   public PyBlock(@Nullable PyBlock parent,
  96                  @NotNull ASTNode node,
  97                  @Nullable Alignment alignment,
  98                  @NotNull Indent indent,
  99                  @Nullable Wrap wrap,
 100                  @NotNull PyBlockContext context) {
 101     myParent = parent;
 102     myAlignment = alignment;
 103     myIndent = indent;
 104     myNode = node;
 105     myWrap = wrap;
 106     myContext = context;
 107     myEmptySequence = isEmptySequence(node);
 108
 109     if (node.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION) {
 110       myDictAlignment = Alignment.createAlignment(true);
 111       myDictWrapping = Wrap.createWrap(myContext.getPySettings().DICT_WRAPPING, true);
 112     }
 113     else {
 114       myDictAlignment = null;
 115       myDictWrapping = null;
 116     }
 117   }
 118
 119   @Override
 120   @NotNull
 121   public ASTNode getNode() {
 122     return myNode;
 123   }
 124
 125   @Override
 126   @NotNull
 127   public TextRange getTextRange() {
 128     return myNode.getTextRange();
 129   }
 130
 131   private Alignment getAlignmentForChildren() {
 132     if (myChildAlignment == null) {
 133       myChildAlignment = Alignment.createAlignment();
 134     }
 135     return myChildAlignment;
 136   }
 137
 138   @Override
 139   @NotNull
 140   public List<Block> getSubBlocks() {
 141     if (mySubBlocks == null) {
 142       mySubBlockByNode = buildSubBlocks();
 143       mySubBlocks = new ArrayList<PyBlock>(mySubBlockByNode.values());
 144     }
 145     return Collections.<Block>unmodifiableList(mySubBlocks);
 146   }
 147
 148   @Nullable
 149   private PyBlock getSubBlockByNode(@NotNull ASTNode node) {
 150     return mySubBlockByNode.get(node);
 151   }
 152
 153   @Nullable
 154   private PyBlock getSubBlockByIndex(int index) {
 155     return mySubBlocks.get(index);
 156   }
 157
 158   @NotNull
 159   private Map<ASTNode, PyBlock> buildSubBlocks() {
 160     final Map<ASTNode, PyBlock> blocks = new LinkedHashMap<ASTNode, PyBlock>();
 161     for (ASTNode child = myNode.getFirstChildNode(); child != null; child = child.getTreeNext()) {
 162
 163       final IElementType childType = child.getElementType();
 164
 165       if (child.getTextRange().isEmpty()) continue;
 166
 167       if (childType == TokenType.WHITE_SPACE) {
 168         continue;
 169       }
 170
 171       blocks.put(child, buildSubBlock(child));
 172     }
 173     return Collections.unmodifiableMap(blocks);
 174   }
 175
 176   @NotNull
 177   private PyBlock buildSubBlock(@NotNull ASTNode child) {
 178     final IElementType parentType = myNode.getElementType();
 179
 180     final ASTNode grandParentNode = myNode.getTreeParent();
 181     final IElementType grandparentType = grandParentNode == null ? null : grandParentNode.getElementType();
 182
 183     final IElementType childType = child.getElementType();
 184     Wrap wrap = null;
 185     Indent childIndent = Indent.getNoneIndent();
 186     Alignment childAlignment = null;
 187
 188     if (parentType == PyElementTypes.BINARY_EXPRESSION && !isInControlStatement()) {
 189       //Setup alignments for binary expression
 190       childAlignment = getAlignmentForChildren();
 191
 192       PyBlock p = myParent; //Check grandparents
 193       while (p != null) {
 194         final ASTNode pNode = p.getNode();
 195         if (ourListElementTypes.contains(pNode.getElementType())) {
 196           if (needListAlignment(child) && !myEmptySequence) {
 197
 198             childAlignment = p.getChildAlignment();
 199             break;
 200           }
 201         }
 202         else if (pNode == PyElementTypes.BINARY_EXPRESSION) {
 203           childAlignment = p.getChildAlignment();
 204         }
 205         if (!breaksAlignment(pNode.getElementType())) {
 206           p = p.myParent;
 207         }
 208         else {
 209           break;
 210         }
 211       }
 212     }
 213
 214     if (childType == PyElementTypes.STATEMENT_LIST) {
 215       if (hasLineBreaksBeforeInSameParent(child, 1) || needLineBreakInStatement()) {
 216         childIndent = Indent.getNormalIndent();
 217       }
 218     }
 219     else if (childType == PyElementTypes.IMPORT_ELEMENT) {
 220       wrap = Wrap.createWrap(WrapType.NORMAL, true);
 221       childIndent = Indent.getNormalIndent();
 222     }
 223     if (childType == PyTokenTypes.END_OF_LINE_COMMENT && parentType == PyElementTypes.FROM_IMPORT_STATEMENT) {
 224       childIndent = Indent.getNormalIndent();
 225     }
 226     if (ourListElementTypes.contains(parentType)) {
 227       // wrapping in non-parenthesized tuple expression is not allowed (PY-1792)
 228       if ((parentType != PyElementTypes.TUPLE_EXPRESSION || grandparentType == PyElementTypes.PARENTHESIZED_EXPRESSION) &&
 229           !ourBrackets.contains(childType) &&
 230           childType != PyTokenTypes.COMMA &&
 231           !isSliceOperand(child) /*&& !isSubscriptionOperand(child)*/) {
 232         wrap = Wrap.createWrap(WrapType.NORMAL, true);
 233       }
 234       if (needListAlignment(child) && !myEmptySequence) {
 235         childAlignment = getAlignmentForChildren();
 236       }
 237       if (childType == PyTokenTypes.END_OF_LINE_COMMENT) {
 238         childIndent = Indent.getNormalIndent();
 239       }
 240     }
 241     else if (parentType == PyElementTypes.BINARY_EXPRESSION &&
 242              (PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens().contains(childType) ||
 243               PyTokenTypes.OPERATIONS.contains(childType))) {
 244       if (isInControlStatement()) {
 245         final PyParenthesizedExpression parens = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyParenthesizedExpression.class, true,
 246                                                                        PyStatementPart.class);
 247         childIndent = parens != null ? Indent.getNormalIndent() : Indent.getContinuationIndent();
 248       }
 249     }
 250
 251     final PyCodeStyleSettings settings = CodeStyleSettingsManager.getSettings(child.getPsi().getProject()).getCustomSettings(PyCodeStyleSettings.class);
 252     if (parentType == PyElementTypes.LIST_LITERAL_EXPRESSION || parentType == PyElementTypes.LIST_COMP_EXPRESSION) {
 253       if (childType == PyTokenTypes.RBRACKET || childType == PyTokenTypes.LBRACKET) {
 254         childIndent = Indent.getNoneIndent();
 255       }
 256       else {
 257         childIndent = Indent.getNormalIndent();
 258       }
 259     }
 260     else if (parentType == PyElementTypes.DICT_LITERAL_EXPRESSION || parentType == PyElementTypes.SET_LITERAL_EXPRESSION ||
 261              parentType == PyElementTypes.SET_COMP_EXPRESSION || parentType == PyElementTypes.DICT_COMP_EXPRESSION) {
 262       if (childType == PyTokenTypes.RBRACE || !hasLineBreaksBeforeInSameParent(child, 1)) {
 263         childIndent = Indent.getNoneIndent();
 264       }
 265       else {
 266         childIndent = Indent.getNormalIndent();
 267       }
 268     }
 269     else if (parentType == PyElementTypes.STRING_LITERAL_EXPRESSION) {
 270       if (PyTokenTypes.STRING_NODES.contains(childType)) {
 271         childAlignment = getAlignmentForChildren();
 272       }
 273     }
 274     else if (parentType == PyElementTypes.FROM_IMPORT_STATEMENT) {
 275       if (myNode.findChildByType(PyTokenTypes.LPAR) != null) {
 276         if (childType == PyElementTypes.IMPORT_ELEMENT) {
 277           if (settings.ALIGN_MULTILINE_IMPORTS) {
 278             childAlignment = getAlignmentForChildren();
 279           }
 280           else {
 281             childIndent = Indent.getNormalIndent();
 282           }
 283         }
 284         if (childType == PyTokenTypes.RPAR) {
 285           childIndent = Indent.getNoneIndent();
 286           if (!hasHangingIndent(myNode.getPsi())) {
 287             childAlignment = getAlignmentForChildren();
 288           }
 289         }
 290       }
 291     }
 292     else if (isValueOfKeyValuePair(child)) {
 293       childIndent = Indent.getNormalIndent();
 294     }
 295     //Align elements vertically if there is an argument in the first line of parenthesized expression
 296     else if (!hasHangingIndent(myNode.getPsi()) &&
 297              ((parentType == PyElementTypes.PARENTHESIZED_EXPRESSION && myContext.getSettings().ALIGN_MULTILINE_PARENTHESIZED_EXPRESSION) ||
 298               (parentType == PyElementTypes.ARGUMENT_LIST && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS) ||
 299               (parentType == PyElementTypes.PARAMETER_LIST && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS)) &&
 300              !isIndentNext(child) &&
 301              !hasLineBreaksBeforeInSameParent(myNode.getFirstChildNode(), 1) &&
 302              !ourListElementTypes.contains(childType)) {
 303
 304       if (!ourBrackets.contains(childType)) {
 305         childAlignment = getAlignmentForChildren();
 306         if (parentType != PyElementTypes.CALL_EXPRESSION) {
 307           childIndent = Indent.getNormalIndent();
 308         }
 309       }
 310       else if (childType == PyTokenTypes.RPAR) {
 311         childIndent = Indent.getNoneIndent();
 312       }
 313     }
 314     else if (parentType == PyElementTypes.GENERATOR_EXPRESSION || parentType == PyElementTypes.PARENTHESIZED_EXPRESSION) {
 315       if (childType == PyTokenTypes.RPAR || !hasLineBreaksBeforeInSameParent(child, 1)) {
 316         childIndent = Indent.getNoneIndent();
 317       }
 318       else {
 319         childIndent = isIndentNext(child) ? Indent.getContinuationIndent() : Indent.getNormalIndent();
 320       }
 321     }
 322     else if (parentType == PyElementTypes.ARGUMENT_LIST || parentType == PyElementTypes.PARAMETER_LIST) {
 323       if (childType == PyTokenTypes.RPAR) {
 324         childIndent = Indent.getNoneIndent();
 325       }
 326       else if (childType != PyTokenTypes.LPAR){
 327         childIndent = Indent.getContinuationIndent();
 328       }
 329     }
 330     else if (parentType == PyElementTypes.SUBSCRIPTION_EXPRESSION) {
 331       final PyExpression indexExpression = ((PySubscriptionExpression)myNode.getPsi()).getIndexExpression();
 332       if (indexExpression != null && child == indexExpression.getNode()) {
 333         childIndent = Indent.getNormalIndent();
 334       }
 335     }
 336     else if (parentType == PyElementTypes.REFERENCE_EXPRESSION) {
 337       if (child != myNode.getFirstChildNode()) {
 338         childIndent = Indent.getNormalIndent();
 339         if (hasLineBreaksBeforeInSameParent(child, 1)) {
 340           if (isInControlStatement()) {
 341             childIndent = Indent.getContinuationIndent();
 342           }
 343           else {
 344             PyBlock b = myParent;
 345             while (b != null) {
 346               if (b.getNode().getPsi() instanceof PyParenthesizedExpression ||
 347                   b.getNode().getPsi() instanceof PyArgumentList ||
 348                   b.getNode().getPsi() instanceof PyParameterList) {
 349                 childAlignment = getAlignmentOfChild(b, 1);
 350                 break;
 351               }
 352               b = b.myParent;
 353             }
 354           }
 355         }
 356       }
 357     }
 358     if (childType == PyElementTypes.KEY_VALUE_EXPRESSION && isChildOfDictLiteral(child)) {
 359       wrap = myDictWrapping;
 360       childIndent = Indent.getNormalIndent();
 361     }
 362
 363     if (isAfterStatementList(child) && !hasLineBreaksBeforeInSameParent(child, 2) && child.getElementType() != PyTokenTypes.END_OF_LINE_COMMENT) {
 364       // maybe enter was pressed and cut us from a previous (nested) statement list
 365       childIndent = Indent.getNormalIndent();
 366     }
 367
 368     if (settings.DICT_ALIGNMENT == DICT_ALIGNMENT_ON_VALUE) {
 369       if (isValueOfKeyValuePairOfDictLiteral(child) && !ourListElementTypes.contains(childType)) {
 370         childAlignment = myParent.myDictAlignment;
 371       }
 372       else if (isValueOfKeyValuePairOfDictLiteral(myNode) &&
 373                ourListElementTypes.contains(parentType) &&
 374                PyTokenTypes.OPEN_BRACES.contains(childType)) {
 375         childAlignment = myParent.myParent.myDictAlignment;
 376       }
 377     }
 378     else if (myContext.getPySettings().DICT_ALIGNMENT == DICT_ALIGNMENT_ON_COLON) {
 379       if (isChildOfKeyValuePairOfDictLiteral(child) && childType == PyTokenTypes.COLON) {
 380         childAlignment = myParent.myDictAlignment;
 381       }
 382     }
 383
 384     ASTNode prev = child.getTreePrev();
 385     while (prev != null && prev.getElementType() == TokenType.WHITE_SPACE) {
 386       if (prev.textContains('\\') &&
 387           !childIndent.equals(Indent.getContinuationIndent(false)) &&
 388           !childIndent.equals(Indent.getContinuationIndent(true))) {
 389         childIndent = isIndentNext(child) ? Indent.getContinuationIndent() : Indent.getNormalIndent();
 390         break;
 391       }
 392       prev = prev.getTreePrev();
 393     }
 394
 395     return new PyBlock(this, child, childAlignment, childIndent, wrap, myContext);
 396   }
 397
 398   private static boolean isValueOfKeyValuePairOfDictLiteral(@NotNull ASTNode node) {
 399     return isValueOfKeyValuePair(node) && isChildOfDictLiteral(node.getTreeParent());
 400   }
 401
 402   private static boolean isChildOfKeyValuePairOfDictLiteral(@NotNull ASTNode node) {
 403     return isChildOfKeyValuePair(node) && isChildOfDictLiteral(node.getTreeParent());
 404   }
 405
 406   private static boolean isChildOfDictLiteral(@NotNull ASTNode node) {
 407     final ASTNode nodeParent = node.getTreeParent();
 408     return nodeParent != null && nodeParent.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION;
 409   }
 410
 411   private static boolean isChildOfKeyValuePair(@NotNull ASTNode node) {
 412     final ASTNode nodeParent = node.getTreeParent();
 413     return nodeParent != null && nodeParent.getElementType() == PyElementTypes.KEY_VALUE_EXPRESSION;
 414   }
 415
 416   private static boolean isValueOfKeyValuePair(@NotNull ASTNode node) {
 417     return isChildOfKeyValuePair(node) && node.getTreeParent().getPsi(PyKeyValueExpression.class).getValue() == node.getPsi();
 418   }
 419
 420   private static boolean isEmptySequence(@NotNull ASTNode node) {
 421     return node.getPsi() instanceof PySequenceExpression && ((PySequenceExpression)node.getPsi()).isEmpty();
 422   }
 423
 424   // Check https://www.python.org/dev/peps/pep-0008/#indentation
 425   private static boolean hasHangingIndent(@NotNull PsiElement elem) {
 426     if (elem instanceof PyCallExpression) {
 427       final PyArgumentList argumentList = ((PyCallExpression)elem).getArgumentList();
 428       return argumentList != null && hasHangingIndent(argumentList);
 429     }
 430     else if (elem instanceof PyFunction) {
 431       return hasHangingIndent(((PyFunction)elem).getParameterList());
 432     }
 433
 434     final PsiElement firstChild;
 435     if (elem instanceof PyFromImportStatement) {
 436       firstChild = ((PyFromImportStatement)elem).getLeftParen();
 437     }
 438     else {
 439       firstChild = elem.getFirstChild();
 440     }
 441     if (firstChild == null) {
 442       return false;
 443     }
 444     final IElementType elementType = elem.getNode().getElementType();
 445     final ASTNode firstChildNode = firstChild.getNode();
 446     if (ourHangingIndentOwners.contains(elementType) && PyTokenTypes.OPEN_BRACES.contains(firstChildNode.getElementType())) {
 447       if (hasLineBreakAfterIgnoringComments(firstChildNode)) {
 448         return true;
 449       }
 450       final PsiElement[] items = getItems(elem);
 451       if (items.length == 0) {
 452         return !PyTokenTypes.CLOSE_BRACES.contains(elem.getLastChild().getNode().getElementType());
 453       }
 454       else {
 455         final PsiElement firstItem = items[0];
 456         if (firstItem instanceof PyNamedParameter) {
 457           final PyExpression defaultValue = ((PyNamedParameter)firstItem).getDefaultValue();
 458           return defaultValue != null && hasHangingIndent(defaultValue);
 459         }
 460         else if (firstItem instanceof PyKeywordArgument) {
 461           final PyExpression valueExpression = ((PyKeywordArgument)firstItem).getValueExpression();
 462           return valueExpression != null && hasHangingIndent(valueExpression);
 463         }
 464         else if (firstItem instanceof PyKeyValueExpression) {
 465           final PyExpression value = ((PyKeyValueExpression)firstItem).getValue();
 466           return value != null && hasHangingIndent(value);
 467         }
 468         return hasHangingIndent(firstItem);
 469       }
 470     }
 471     else {
 472       return false;
 473     }
 474   }
 475
 476   @NotNull
 477   private static PsiElement[] getItems(@NotNull PsiElement elem) {
 478     if (elem instanceof PySequenceExpression) {
 479       return ((PySequenceExpression)elem).getElements();
 480     }
 481     else if (elem instanceof PyParameterList) {
 482       return ((PyParameterList)elem).getParameters();
 483     }
 484     else if (elem instanceof PyArgumentList) {
 485       return ((PyArgumentList)elem).getArguments();
 486     }
 487     else if (elem instanceof PyFromImportStatement) {
 488       return ((PyFromImportStatement)elem).getImportElements();
 489     }
 490     else if (elem instanceof PyParenthesizedExpression) {
 491       final PyExpression containedExpression = ((PyParenthesizedExpression)elem).getContainedExpression();
 492       if (containedExpression instanceof PyTupleExpression) {
 493         return ((PyTupleExpression)containedExpression).getElements();
 494       }
 495       else if (containedExpression != null) {
 496         return new PsiElement[]{containedExpression};
 497       }
 498     }
 499     return PsiElement.EMPTY_ARRAY;
 500   }
 501
 502   private static boolean breaksAlignment(IElementType type) {
 503     return type != PyElementTypes.BINARY_EXPRESSION;
 504   }
 505
 506   private static Alignment getAlignmentOfChild(@NotNull PyBlock b, int childNum) {
 507     if (b.getSubBlocks().size() > childNum) {
 508       final ChildAttributes attributes = b.getChildAttributes(childNum);
 509       return attributes.getAlignment();
 510     }
 511     return null;
 512   }
 513
 514   private static boolean isIndentNext(@NotNull ASTNode child) {
 515     final PsiElement psi = PsiTreeUtil.getParentOfType(child.getPsi(), PyStatement.class);
 516
 517     return psi instanceof PyIfStatement ||
 518            psi instanceof PyForStatement ||
 519            psi instanceof PyWithStatement ||
 520            psi instanceof PyClass ||
 521            psi instanceof PyFunction ||
 522            psi instanceof PyTryExceptStatement ||
 523            psi instanceof PyElsePart ||
 524            psi instanceof PyIfPart ||
 525            psi instanceof PyWhileStatement;
 526   }
 527
 528   private static boolean isSubscriptionOperand(@NotNull ASTNode child) {
 529     return child.getTreeParent().getElementType() == PyElementTypes.SUBSCRIPTION_EXPRESSION &&
 530            child.getPsi() == ((PySubscriptionExpression)child.getTreeParent().getPsi()).getOperand();
 531   }
 532
 533   private boolean isInControlStatement() {
 534     return getControlStatementHeader(myNode) != null;
 535   }
 536
 537   @Nullable
 538   private static PsiElement getControlStatementHeader(@NotNull ASTNode node) {
 539     final PyStatementPart statementPart = PsiTreeUtil.getParentOfType(node.getPsi(), PyStatementPart.class, false, PyStatementList.class);
 540     if (statementPart != null) {
 541       return statementPart;
 542     }
 543     final PyWithItem withItem = PsiTreeUtil.getParentOfType(node.getPsi(), PyWithItem.class);
 544     if (withItem != null) {
 545       return withItem.getParent();
 546     }
 547     return null;
 548   }
 549
 550   private boolean isSliceOperand(@NotNull ASTNode child) {
 551     if (myNode.getPsi() instanceof PySliceExpression) {
 552       final PySliceExpression sliceExpression = (PySliceExpression)myNode.getPsi();
 553       final PyExpression operand = sliceExpression.getOperand();
 554       return operand.getNode() == child;
 555     }
 556     return false;
 557   }
 558
 559   private static boolean isAfterStatementList(@NotNull ASTNode child) {
 560     final PsiElement prev = child.getPsi().getPrevSibling();
 561     if (!(prev instanceof PyStatement)) {
 562       return false;
 563     }
 564     final PsiElement lastChild = PsiTreeUtil.getDeepestLast(prev);
 565     return lastChild.getParent() instanceof PyStatementList;
 566   }
 567
 568   private boolean needListAlignment(@NotNull ASTNode child) {
 569     final IElementType childType = child.getElementType();
 570     if (PyTokenTypes.OPEN_BRACES.contains(childType)) {
 571       return false;
 572     }
 573     if (PyTokenTypes.CLOSE_BRACES.contains(childType)) {
 574       final ASTNode prevNonSpace = findPrevNonSpaceNode(child);
 575       if (prevNonSpace != null &&
 576           prevNonSpace.getElementType() == PyTokenTypes.COMMA &&
 577           myContext.getMode() == FormattingMode.ADJUST_INDENT) {
 578         return true;
 579       }
 580       return !hasHangingIndent(myNode.getPsi());
 581     }
 582     if (myNode.getElementType() == PyElementTypes.ARGUMENT_LIST) {
 583       if (!myContext.getSettings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS || hasHangingIndent(myNode.getPsi())) {
 584         return false;
 585       }
 586       if (child.getElementType() == PyTokenTypes.COMMA) {
 587         return false;
 588       }
 589       return true;
 590     }
 591     if (myNode.getElementType() == PyElementTypes.PARAMETER_LIST) {
 592       return !hasHangingIndent(myNode.getPsi()) && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS;
 593     }
 594     if (myNode.getElementType() == PyElementTypes.SUBSCRIPTION_EXPRESSION) {
 595       return false;
 596     }
 597     if (child.getElementType() == PyTokenTypes.COMMA) {
 598       return false;
 599     }
 600     return myContext.getPySettings().ALIGN_COLLECTIONS_AND_COMPREHENSIONS && !hasHangingIndent(myNode.getPsi());
 601   }
 602
 603   @Nullable
 604   private static ASTNode findPrevNonSpaceNode(@NotNull ASTNode node) {
 605     do {
 606       node = node.getTreePrev();
 607     }
 608     while (isWhitespace(node));
 609     return node;
 610   }
 611
 612   private static boolean isWhitespace(@Nullable ASTNode node) {
 613     return node != null && (node.getElementType() == TokenType.WHITE_SPACE || PyTokenTypes.WHITESPACE.contains(node.getElementType()));
 614   }
 615
 616   private static boolean hasLineBreaksBeforeInSameParent(@NotNull ASTNode node, int minCount) {
 617     final ASTNode treePrev = node.getTreePrev();
 618     return (treePrev != null && isWhitespaceWithLineBreaks(TreeUtil.findLastLeaf(treePrev), minCount)) ||
 619            // Can happen, e.g. when you delete a statement from the beginning of a statement list
 620            isWhitespaceWithLineBreaks(node.getFirstChildNode(), minCount);
 621   }
 622
 623   private static boolean hasLineBreakAfterIgnoringComments(@NotNull ASTNode node) {
 624     for (ASTNode next = TreeUtil.nextLeaf(node); next != null; next = TreeUtil.nextLeaf(next)) {
 625       if (isWhitespace(next)) {
 626         if (next.textContains('\n')) {
 627           return true;
 628         }
 629       }
 630       else if (next.getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
 631         return true;
 632       }
 633       else {
 634         break;
 635       }
 636     }
 637     return false;
 638   }
 639
 640   private static boolean isWhitespaceWithLineBreaks(@Nullable ASTNode node, int minCount) {
 641     if (isWhitespace(node)) {
 642       if (minCount == 1) {
 643         return node.textContains('\n');
 644       }
 645       final String prevNodeText = node.getText();
 646       int count = 0;
 647       for (int i = 0; i < prevNodeText.length(); i++) {
 648         if (prevNodeText.charAt(i) == '\n') {
 649           count++;
 650           if (count == minCount) {
 651             return true;
 652           }
 653         }
 654       }
 655     }
 656     return false;
 657   }
 658
 659   @Override
 660   @Nullable
 661   public Wrap getWrap() {
 662     return myWrap;
 663   }
 664
 665   @Override
 666   @Nullable
 667   public Indent getIndent() {
 668     assert myIndent != null;
 669     return myIndent;
 670   }
 671
 672   @Override
 673   @Nullable
 674   public Alignment getAlignment() {
 675     return myAlignment;
 676   }
 677
 678   @Override
 679   @Nullable
 680   public Spacing getSpacing(@Nullable Block child1, @NotNull Block child2) {
 681     if (child1 instanceof ASTBlock && child2 instanceof ASTBlock) {
 682       final ASTNode node1 = ((ASTBlock)child1).getNode();
 683       ASTNode node2 = ((ASTBlock)child2).getNode();
 684       final IElementType childType1 = node1.getElementType();
 685       final PsiElement psi1 = node1.getPsi();
 686
 687       PsiElement psi2 = node2.getPsi();
 688       // skip not inline comments to handles blank lines between various declarations
 689       if (psi2 instanceof PsiComment && hasLineBreaksBeforeInSameParent(node2, 1)) {
 690         final PsiElement nonCommentAfter = PyPsiUtils.getNextNonCommentSibling(psi2, true);
 691         if (nonCommentAfter != null) {
 692           psi2 = nonCommentAfter;
 693         }
 694       }
 695       node2 = psi2.getNode();
 696       final IElementType childType2 = psi2.getNode().getElementType();
 697       //noinspection ConstantConditions
 698       child2 = getSubBlockByNode(node2);
 699       final CommonCodeStyleSettings settings = myContext.getSettings();
 700
 701       if ((childType1 == PyTokenTypes.EQ || childType2 == PyTokenTypes.EQ)) {
 702         final PyNamedParameter namedParameter = as(myNode.getPsi(), PyNamedParameter.class);
 703         if (namedParameter != null && namedParameter.getAnnotation() != null) {
 704           return Spacing.createSpacing(1, 1, 0, settings.KEEP_LINE_BREAKS, settings.KEEP_BLANK_LINES_IN_CODE);
 705         }
 706       }
 707
 708       if (childType1 == PyTokenTypes.COLON && psi2 instanceof PyStatementList) {
 709         if (needLineBreakInStatement()) {
 710           return Spacing.createSpacing(0, 0, 1, true, settings.KEEP_BLANK_LINES_IN_CODE);
 711         }
 712       }
 713
 714       if ((PyElementTypes.CLASS_OR_FUNCTION.contains(childType1) && STATEMENT_OR_DECLARATION.contains(childType2)) ||
 715           STATEMENT_OR_DECLARATION.contains(childType1) && PyElementTypes.CLASS_OR_FUNCTION.contains(childType2)) {
 716         if (PyUtil.isTopLevel(psi1)) {
 717           return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS);
 718         }
 719       }
 720
 721       if (psi1 instanceof PyImportStatementBase) {
 722         if (psi2 instanceof PyImportStatementBase) {
 723           final Boolean leftImportIsGroupStart = psi1.getCopyableUserData(IMPORT_GROUP_BEGIN);
 724           final Boolean rightImportIsGroupStart = psi2.getCopyableUserData(IMPORT_GROUP_BEGIN);
 725           // Cleanup user data, it's no longer needed
 726           psi1.putCopyableUserData(IMPORT_GROUP_BEGIN, null);
 727           // Don't remove IMPORT_GROUP_BEGIN from the element psi2 yet, because spacing is constructed pairwise:
 728           // it might be needed on the next iteration.
 729           //psi2.putCopyableUserData(IMPORT_GROUP_BEGIN, null);
 730           if (rightImportIsGroupStart != null) {
 731             return Spacing.createSpacing(0, 0, 2, true, 1);
 732           }
 733           else if (leftImportIsGroupStart != null) {
 734             // It's a trick to keep spacing consistent when new import statement is inserted
 735             // at the beginning of an import group, i.e. if there is a blank line before the next
 736             // import we want to save it, but remove line *after* inserted import.
 737             return Spacing.createSpacing(0, 0, 1, false, 0);
 738           }
 739         }
 740         if (psi2 instanceof PyStatement && !(psi2 instanceof PyImportStatementBase)) {
 741           if (PyUtil.isTopLevel(psi1)) {
 742             return getBlankLinesForOption(settings.BLANK_LINES_AFTER_IMPORTS);
 743           }
 744           else {
 745             return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AFTER_LOCAL_IMPORTS);
 746           }
 747         }
 748       }
 749
 750       if (psi2 instanceof PsiComment && !hasLineBreaksBeforeInSameParent(psi2.getNode(), 1) && myContext.getPySettings().SPACE_BEFORE_NUMBER_SIGN) {
 751         return Spacing.createSpacing(2, 0, 0, false, 0);
 752       }
 753     }
 754     return myContext.getSpacingBuilder().getSpacing(this, child1, child2);
 755   }
 756
 757   @NotNull
 758   private Spacing getBlankLinesForOption(int option) {
 759     final int blankLines = option + 1;
 760     return Spacing.createSpacing(0, 0, blankLines,
 761                                  myContext.getSettings().KEEP_LINE_BREAKS,
 762                                  myContext.getSettings().KEEP_BLANK_LINES_IN_DECLARATIONS);
 763   }
 764
 765   private boolean needLineBreakInStatement() {
 766     final PyStatement statement = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyStatement.class);
 767     if (statement != null) {
 768       final Collection<PyStatementPart> parts = PsiTreeUtil.collectElementsOfType(statement, PyStatementPart.class);
 769       return (parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
 770              (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE);
 771     }
 772     return false;
 773   }
 774
 775   @Override
 776   @NotNull
 777   public ChildAttributes getChildAttributes(int newChildIndex) {
 778     int statementListsBelow = 0;
 779     if (newChildIndex > 0) {
 780       // always pass decision to a sane block from top level from file or definition
 781       if (myNode.getPsi() instanceof PyFile || myNode.getElementType() == PyTokenTypes.COLON) {
 782         return ChildAttributes.DELEGATE_TO_PREV_CHILD;
 783       }
 784
 785       final PyBlock insertAfterBlock = getSubBlockByIndex(newChildIndex - 1);
 786
 787       final ASTNode prevNode = insertAfterBlock.getNode();
 788       final PsiElement prevElt = prevNode.getPsi();
 789
 790       // stmt lists, parts and definitions should also think for themselves
 791       if (prevElt instanceof PyStatementList) {
 792         if (dedentAfterLastStatement((PyStatementList)prevElt)) {
 793           return new ChildAttributes(Indent.getNoneIndent(), getChildAlignment());
 794         }
 795         return ChildAttributes.DELEGATE_TO_PREV_CHILD;
 796       }
 797       else if (prevElt instanceof PyStatementPart) {
 798         return ChildAttributes.DELEGATE_TO_PREV_CHILD;
 799       }
 800
 801       ASTNode lastChild = insertAfterBlock.getNode();
 802
 803       // HACK? This code fragment is needed to make testClass2() pass,
 804       // but I don't quite understand why it is necessary and why the formatter
 805       // doesn't request childAttributes from the correct block
 806       while (lastChild != null) {
 807         final IElementType lastType = lastChild.getElementType();
 808         if (lastType == PyElementTypes.STATEMENT_LIST && hasLineBreaksBeforeInSameParent(lastChild, 1)) {
 809           if (dedentAfterLastStatement((PyStatementList)lastChild.getPsi())) {
 810             break;
 811           }
 812           statementListsBelow++;
 813         }
 814         else if (statementListsBelow > 0 && lastChild.getPsi() instanceof PsiErrorElement) {
 815           statementListsBelow++;
 816         }
 817         if (myNode.getElementType() == PyElementTypes.STATEMENT_LIST && lastChild.getPsi() instanceof PsiErrorElement) {
 818           return ChildAttributes.DELEGATE_TO_PREV_CHILD;
 819         }
 820         lastChild = getLastNonSpaceChild(lastChild, true);
 821       }
 822     }
 823
 824     // HACKETY-HACK
 825     // If a multi-step dedent follows the cursor position (see testMultiDedent()),
 826     // the whitespace (which must be a single Py:LINE_BREAK token) gets attached
 827     // to the outermost indented block (because we may not consume the DEDENT
 828     // tokens while parsing inner blocks). The policy is to put the indent to
 829     // the innermost block, so we need to resolve the situation here. Nested
 830     // delegation sometimes causes NPEs in formatter core, so we calculate the
 831     // correct indent manually.
 832     if (statementListsBelow > 0) { // was 1... strange
 833       @SuppressWarnings("ConstantConditions") final
 834       int indent = myContext.getSettings().getIndentOptions().INDENT_SIZE;
 835       return new ChildAttributes(Indent.getSpaceIndent(indent * statementListsBelow), null);
 836     }
 837
 838     /*
 839     // it might be something like "def foo(): # comment" or "[1, # comment"; jump up to the real thing
 840     if (_node instanceof PsiComment || _node instanceof PsiWhiteSpace) {
 841       get
 842     }
 843     */
 844
 845
 846     final Indent childIndent = getChildIndent(newChildIndex);
 847     final Alignment childAlignment = getChildAlignment();
 848     return new ChildAttributes(childIndent, childAlignment);
 849   }
 850
 851   private static boolean dedentAfterLastStatement(@NotNull PyStatementList statementList) {
 852     final PyStatement[] statements = statementList.getStatements();
 853     if (statements.length == 0) {
 854       return false;
 855     }
 856     final PyStatement last = statements[statements.length - 1];
 857     return last instanceof PyReturnStatement || last instanceof PyRaiseStatement || last instanceof PyPassStatement;
 858   }
 859
 860   @Nullable
 861   private Alignment getChildAlignment() {
 862     if (ourListElementTypes.contains(myNode.getElementType())) {
 863       if (isInControlStatement()) {
 864         return null;
 865       }
 866       if (myNode.getPsi() instanceof PyParameterList && !myContext.getSettings().ALIGN_MULTILINE_PARAMETERS) {
 867         return null;
 868       }
 869       if (myNode.getPsi() instanceof PyDictLiteralExpression) {
 870         final PyKeyValueExpression lastElement = ArrayUtil.getLastElement(((PyDictLiteralExpression)myNode.getPsi()).getElements());
 871         if (lastElement == null || lastElement.getValue() == null /* incomplete */) {
 872           return null;
 873         }
 874       }
 875       return getAlignmentForChildren();
 876     }
 877     return null;
 878   }
 879
 880   @NotNull
 881   private Indent getChildIndent(int newChildIndex) {
 882     final ASTNode afterNode = getAfterNode(newChildIndex);
 883     final ASTNode lastChild = getLastNonSpaceChild(myNode, false);
 884     if (lastChild != null && lastChild.getElementType() == PyElementTypes.STATEMENT_LIST && mySubBlocks.size() >= newChildIndex) {
 885       if (afterNode == null) {
 886         return Indent.getNoneIndent();
 887       }
 888
 889       // handle pressing Enter after colon and before first statement in
 890       // existing statement list
 891       if (afterNode.getElementType() == PyElementTypes.STATEMENT_LIST || afterNode.getElementType() == PyTokenTypes.COLON) {
 892         return Indent.getNormalIndent();
 893       }
 894
 895       // handle pressing Enter after colon when there is nothing in the
 896       // statement list
 897       final ASTNode lastFirstChild = lastChild.getFirstChildNode();
 898       if (lastFirstChild != null && lastFirstChild == lastChild.getLastChildNode() && lastFirstChild.getPsi() instanceof PsiErrorElement) {
 899         return Indent.getNormalIndent();
 900       }
 901     }
 902     else if (lastChild != null && PyElementTypes.LIST_LIKE_EXPRESSIONS.contains(lastChild.getElementType())) {
 903       // handle pressing enter at the end of a list literal when there's no closing paren or bracket
 904       final ASTNode lastLastChild = lastChild.getLastChildNode();
 905       if (lastLastChild != null && lastLastChild.getPsi() instanceof PsiErrorElement) {
 906         // we're at a place like this: [foo, ... bar, <caret>
 907         // we'd rather align to foo. this may be not a multiple of tabs.
 908         final PsiElement expr = lastChild.getPsi();
 909         PsiElement exprItem = expr.getFirstChild();
 910         boolean found = false;
 911         while (exprItem != null) { // find a worthy element to align to
 912           if (exprItem instanceof PyElement) {
 913             found = true; // align to foo in "[foo,"
 914             break;
 915           }
 916           if (exprItem instanceof PsiComment) {
 917             found = true; // align to foo in "[ # foo,"
 918             break;
 919           }
 920           exprItem = exprItem.getNextSibling();
 921         }
 922         if (found) {
 923           final PsiDocumentManager docMgr = PsiDocumentManager.getInstance(exprItem.getProject());
 924           final Document doc = docMgr.getDocument(exprItem.getContainingFile());
 925           if (doc != null) {
 926             int lineNum = doc.getLineNumber(exprItem.getTextOffset());
 927             final int itemCol = exprItem.getTextOffset() - doc.getLineStartOffset(lineNum);
 928             final PsiElement hereElt = getNode().getPsi();
 929             lineNum = doc.getLineNumber(hereElt.getTextOffset());
 930             final int nodeCol = hereElt.getTextOffset() - doc.getLineStartOffset(lineNum);
 931             final int padding = itemCol - nodeCol;
 932             if (padding > 0) { // negative is a syntax error,  but possible
 933               return Indent.getSpaceIndent(padding);
 934             }
 935           }
 936         }
 937         return Indent.getContinuationIndent(); // a fallback
 938       }
 939     }
 940
 941     if (afterNode != null && afterNode.getElementType() == PyElementTypes.KEY_VALUE_EXPRESSION) {
 942       final PyKeyValueExpression keyValue = (PyKeyValueExpression)afterNode.getPsi();
 943       if (keyValue != null && keyValue.getValue() == null) {  // incomplete
 944         return Indent.getContinuationIndent();
 945       }
 946     }
 947
 948     // constructs that imply indent for their children
 949     if (myNode.getElementType().equals(PyElementTypes.PARAMETER_LIST)) {
 950       return Indent.getContinuationIndent();
 951     }
 952     if (ourListElementTypes.contains(myNode.getElementType()) || myNode.getPsi() instanceof PyStatementPart) {
 953       return Indent.getNormalIndent();
 954     }
 955
 956     if (afterNode != null) {
 957       ASTNode wsAfter = afterNode.getTreeNext();
 958       while (wsAfter != null && wsAfter.getElementType() == TokenType.WHITE_SPACE) {
 959         if (wsAfter.getText().indexOf('\\') >= 0) {
 960           return Indent.getNormalIndent();
 961         }
 962         wsAfter = wsAfter.getTreeNext();
 963       }
 964     }
 965     return Indent.getNoneIndent();
 966   }
 967
 968   @Nullable
 969   private ASTNode getAfterNode(int newChildIndex) {
 970     if (newChildIndex == 0) {  // block text contains backslash line wrappings, child block list not built
 971       return null;
 972     }
 973     int prevIndex = newChildIndex - 1;
 974     while (prevIndex > 0 && getSubBlockByIndex(prevIndex).getNode().getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
 975       prevIndex--;
 976     }
 977     return getSubBlockByIndex(prevIndex).getNode();
 978   }
 979
 980   private static ASTNode getLastNonSpaceChild(@NotNull ASTNode node, boolean acceptError) {
 981     ASTNode lastChild = node.getLastChildNode();
 982     while (lastChild != null &&
 983            (lastChild.getElementType() == TokenType.WHITE_SPACE || (!acceptError && lastChild.getPsi() instanceof PsiErrorElement))) {
 984       lastChild = lastChild.getTreePrev();
 985     }
 986     return lastChild;
 987   }
 988
 989   @Override
 990   public boolean isIncomplete() {
 991     // if there's something following us, we're not incomplete
 992     if (!PsiTreeUtil.hasErrorElements(myNode.getPsi())) {
 993       PsiElement element = myNode.getPsi().getNextSibling();
 994       while (element instanceof PsiWhiteSpace) {
 995         element = element.getNextSibling();
 996       }
 997       if (element != null) {
 998         return false;
 999       }
1000     }
1001
1002     final ASTNode lastChild = getLastNonSpaceChild(myNode, false);
1003     if (lastChild != null) {
1004       if (lastChild.getElementType() == PyElementTypes.STATEMENT_LIST) {
1005         // only multiline statement lists are considered incomplete
1006         final ASTNode statementListPrev = lastChild.getTreePrev();
1007         if (statementListPrev != null && statementListPrev.getText().indexOf('\n') >= 0) {
1008           return true;
1009         }
1010       }
1011       if (lastChild.getElementType() == PyElementTypes.BINARY_EXPRESSION) {
1012         final PyBinaryExpression binaryExpression = (PyBinaryExpression)lastChild.getPsi();
1013         if (binaryExpression.getRightExpression() == null) {
1014           return true;
1015         }
1016       }
1017       if (isIncompleteCall(lastChild)) return true;
1018     }
1019
1020     if (myNode.getPsi() instanceof PyArgumentList) {
1021       final PyArgumentList argumentList = (PyArgumentList)myNode.getPsi();
1022       return argumentList.getClosingParen() == null;
1023     }
1024     if (isIncompleteCall(myNode)) {
1025       return true;
1026     }
1027
1028     return false;
1029   }
1030
1031   private static boolean isIncompleteCall(@NotNull ASTNode node) {
1032     if (node.getElementType() == PyElementTypes.CALL_EXPRESSION) {
1033       final PyCallExpression callExpression = (PyCallExpression)node.getPsi();
1034       final PyArgumentList argumentList = callExpression.getArgumentList();
1035       if (argumentList == null || argumentList.getClosingParen() == null) {
1036         return true;
1037       }
1038     }
1039     return false;
1040   }
1041
1042   @Override
1043   public boolean isLeaf() {
1044     return myNode.getFirstChildNode() == null;
1045   }
1046 }