PY-17564 Keep blank lines between imports inserted by user
[idea/community.git] / 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           if (psi2.getCopyableUserData(IMPORT_GROUP_BEGIN) != null) {
724             return Spacing.createSpacing(0, 0, 2, true, 1);
725           }
726           else if (psi1.getCopyableUserData(IMPORT_GROUP_BEGIN) != null) {
727             // It's a trick to keep spacing consistent when new import statement is inserted
728             // at the beginning of an import group, i.e. if there is a blank line before the next
729             // import we want to save it, but remove line *after* inserted import.
730             return Spacing.createSpacing(0, 0, 1, false, 0);
731           }
732         }
733         if (psi2 instanceof PyStatement && !(psi2 instanceof PyImportStatementBase)) {
734           if (PyUtil.isTopLevel(psi1)) {
735             return getBlankLinesForOption(settings.BLANK_LINES_AFTER_IMPORTS);
736           }
737           else {
738             return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AFTER_LOCAL_IMPORTS);
739           }
740         }
741       }
742
743       if (psi2 instanceof PsiComment && !hasLineBreaksBeforeInSameParent(psi2.getNode(), 1) && myContext.getPySettings().SPACE_BEFORE_NUMBER_SIGN) {
744         return Spacing.createSpacing(2, 0, 0, false, 0);
745       }
746     }
747     return myContext.getSpacingBuilder().getSpacing(this, child1, child2);
748   }
749
750   @NotNull
751   private Spacing getBlankLinesForOption(int option) {
752     final int blankLines = option + 1;
753     return Spacing.createSpacing(0, 0, blankLines,
754                                  myContext.getSettings().KEEP_LINE_BREAKS,
755                                  myContext.getSettings().KEEP_BLANK_LINES_IN_DECLARATIONS);
756   }
757
758   private boolean needLineBreakInStatement() {
759     final PyStatement statement = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyStatement.class);
760     if (statement != null) {
761       final Collection<PyStatementPart> parts = PsiTreeUtil.collectElementsOfType(statement, PyStatementPart.class);
762       return (parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
763              (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE);
764     }
765     return false;
766   }
767
768   @Override
769   @NotNull
770   public ChildAttributes getChildAttributes(int newChildIndex) {
771     int statementListsBelow = 0;
772     if (newChildIndex > 0) {
773       // always pass decision to a sane block from top level from file or definition
774       if (myNode.getPsi() instanceof PyFile || myNode.getElementType() == PyTokenTypes.COLON) {
775         return ChildAttributes.DELEGATE_TO_PREV_CHILD;
776       }
777
778       final PyBlock insertAfterBlock = getSubBlockByIndex(newChildIndex - 1);
779
780       final ASTNode prevNode = insertAfterBlock.getNode();
781       final PsiElement prevElt = prevNode.getPsi();
782
783       // stmt lists, parts and definitions should also think for themselves
784       if (prevElt instanceof PyStatementList) {
785         if (dedentAfterLastStatement((PyStatementList)prevElt)) {
786           return new ChildAttributes(Indent.getNoneIndent(), getChildAlignment());
787         }
788         return ChildAttributes.DELEGATE_TO_PREV_CHILD;
789       }
790       else if (prevElt instanceof PyStatementPart) {
791         return ChildAttributes.DELEGATE_TO_PREV_CHILD;
792       }
793
794       ASTNode lastChild = insertAfterBlock.getNode();
795
796       // HACK? This code fragment is needed to make testClass2() pass,
797       // but I don't quite understand why it is necessary and why the formatter
798       // doesn't request childAttributes from the correct block
799       while (lastChild != null) {
800         final IElementType lastType = lastChild.getElementType();
801         if (lastType == PyElementTypes.STATEMENT_LIST && hasLineBreaksBeforeInSameParent(lastChild, 1)) {
802           if (dedentAfterLastStatement((PyStatementList)lastChild.getPsi())) {
803             break;
804           }
805           statementListsBelow++;
806         }
807         else if (statementListsBelow > 0 && lastChild.getPsi() instanceof PsiErrorElement) {
808           statementListsBelow++;
809         }
810         if (myNode.getElementType() == PyElementTypes.STATEMENT_LIST && lastChild.getPsi() instanceof PsiErrorElement) {
811           return ChildAttributes.DELEGATE_TO_PREV_CHILD;
812         }
813         lastChild = getLastNonSpaceChild(lastChild, true);
814       }
815     }
816
817     // HACKETY-HACK
818     // If a multi-step dedent follows the cursor position (see testMultiDedent()),
819     // the whitespace (which must be a single Py:LINE_BREAK token) gets attached
820     // to the outermost indented block (because we may not consume the DEDENT
821     // tokens while parsing inner blocks). The policy is to put the indent to
822     // the innermost block, so we need to resolve the situation here. Nested
823     // delegation sometimes causes NPEs in formatter core, so we calculate the
824     // correct indent manually.
825     if (statementListsBelow > 0) { // was 1... strange
826       @SuppressWarnings("ConstantConditions") final
827       int indent = myContext.getSettings().getIndentOptions().INDENT_SIZE;
828       return new ChildAttributes(Indent.getSpaceIndent(indent * statementListsBelow), null);
829     }
830
831     /*
832     // it might be something like "def foo(): # comment" or "[1, # comment"; jump up to the real thing
833     if (_node instanceof PsiComment || _node instanceof PsiWhiteSpace) {
834       get
835     }
836     */
837
838
839     final Indent childIndent = getChildIndent(newChildIndex);
840     final Alignment childAlignment = getChildAlignment();
841     return new ChildAttributes(childIndent, childAlignment);
842   }
843
844   private static boolean dedentAfterLastStatement(@NotNull PyStatementList statementList) {
845     final PyStatement[] statements = statementList.getStatements();
846     if (statements.length == 0) {
847       return false;
848     }
849     final PyStatement last = statements[statements.length - 1];
850     return last instanceof PyReturnStatement || last instanceof PyRaiseStatement || last instanceof PyPassStatement;
851   }
852
853   @Nullable
854   private Alignment getChildAlignment() {
855     if (ourListElementTypes.contains(myNode.getElementType())) {
856       if (isInControlStatement()) {
857         return null;
858       }
859       if (myNode.getPsi() instanceof PyParameterList && !myContext.getSettings().ALIGN_MULTILINE_PARAMETERS) {
860         return null;
861       }
862       if (myNode.getPsi() instanceof PyDictLiteralExpression) {
863         final PyKeyValueExpression lastElement = ArrayUtil.getLastElement(((PyDictLiteralExpression)myNode.getPsi()).getElements());
864         if (lastElement == null || lastElement.getValue() == null /* incomplete */) {
865           return null;
866         }
867       }
868       return getAlignmentForChildren();
869     }
870     return null;
871   }
872
873   @NotNull
874   private Indent getChildIndent(int newChildIndex) {
875     final ASTNode afterNode = getAfterNode(newChildIndex);
876     final ASTNode lastChild = getLastNonSpaceChild(myNode, false);
877     if (lastChild != null && lastChild.getElementType() == PyElementTypes.STATEMENT_LIST && mySubBlocks.size() >= newChildIndex) {
878       if (afterNode == null) {
879         return Indent.getNoneIndent();
880       }
881
882       // handle pressing Enter after colon and before first statement in
883       // existing statement list
884       if (afterNode.getElementType() == PyElementTypes.STATEMENT_LIST || afterNode.getElementType() == PyTokenTypes.COLON) {
885         return Indent.getNormalIndent();
886       }
887
888       // handle pressing Enter after colon when there is nothing in the
889       // statement list
890       final ASTNode lastFirstChild = lastChild.getFirstChildNode();
891       if (lastFirstChild != null && lastFirstChild == lastChild.getLastChildNode() && lastFirstChild.getPsi() instanceof PsiErrorElement) {
892         return Indent.getNormalIndent();
893       }
894     }
895     else if (lastChild != null && PyElementTypes.LIST_LIKE_EXPRESSIONS.contains(lastChild.getElementType())) {
896       // handle pressing enter at the end of a list literal when there's no closing paren or bracket 
897       final ASTNode lastLastChild = lastChild.getLastChildNode();
898       if (lastLastChild != null && lastLastChild.getPsi() instanceof PsiErrorElement) {
899         // we're at a place like this: [foo, ... bar, <caret>
900         // we'd rather align to foo. this may be not a multiple of tabs.
901         final PsiElement expr = lastChild.getPsi();
902         PsiElement exprItem = expr.getFirstChild();
903         boolean found = false;
904         while (exprItem != null) { // find a worthy element to align to
905           if (exprItem instanceof PyElement) {
906             found = true; // align to foo in "[foo,"
907             break;
908           }
909           if (exprItem instanceof PsiComment) {
910             found = true; // align to foo in "[ # foo,"
911             break;
912           }
913           exprItem = exprItem.getNextSibling();
914         }
915         if (found) {
916           final PsiDocumentManager docMgr = PsiDocumentManager.getInstance(exprItem.getProject());
917           final Document doc = docMgr.getDocument(exprItem.getContainingFile());
918           if (doc != null) {
919             int lineNum = doc.getLineNumber(exprItem.getTextOffset());
920             final int itemCol = exprItem.getTextOffset() - doc.getLineStartOffset(lineNum);
921             final PsiElement hereElt = getNode().getPsi();
922             lineNum = doc.getLineNumber(hereElt.getTextOffset());
923             final int nodeCol = hereElt.getTextOffset() - doc.getLineStartOffset(lineNum);
924             final int padding = itemCol - nodeCol;
925             if (padding > 0) { // negative is a syntax error,  but possible
926               return Indent.getSpaceIndent(padding);
927             }
928           }
929         }
930         return Indent.getContinuationIndent(); // a fallback
931       }
932     }
933
934     if (afterNode != null && afterNode.getElementType() == PyElementTypes.KEY_VALUE_EXPRESSION) {
935       final PyKeyValueExpression keyValue = (PyKeyValueExpression)afterNode.getPsi();
936       if (keyValue != null && keyValue.getValue() == null) {  // incomplete
937         return Indent.getContinuationIndent();
938       }
939     }
940
941     // constructs that imply indent for their children
942     if (myNode.getElementType().equals(PyElementTypes.PARAMETER_LIST)) {
943       return Indent.getContinuationIndent();
944     }
945     if (ourListElementTypes.contains(myNode.getElementType()) || myNode.getPsi() instanceof PyStatementPart) {
946       return Indent.getNormalIndent();
947     }
948
949     if (afterNode != null) {
950       ASTNode wsAfter = afterNode.getTreeNext();
951       while (wsAfter != null && wsAfter.getElementType() == TokenType.WHITE_SPACE) {
952         if (wsAfter.getText().indexOf('\\') >= 0) {
953           return Indent.getNormalIndent();
954         }
955         wsAfter = wsAfter.getTreeNext();
956       }
957     }
958     return Indent.getNoneIndent();
959   }
960
961   @Nullable
962   private ASTNode getAfterNode(int newChildIndex) {
963     if (newChildIndex == 0) {  // block text contains backslash line wrappings, child block list not built
964       return null;
965     }
966     int prevIndex = newChildIndex - 1;
967     while (prevIndex > 0 && getSubBlockByIndex(prevIndex).getNode().getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
968       prevIndex--;
969     }
970     return getSubBlockByIndex(prevIndex).getNode();
971   }
972
973   private static ASTNode getLastNonSpaceChild(@NotNull ASTNode node, boolean acceptError) {
974     ASTNode lastChild = node.getLastChildNode();
975     while (lastChild != null &&
976            (lastChild.getElementType() == TokenType.WHITE_SPACE || (!acceptError && lastChild.getPsi() instanceof PsiErrorElement))) {
977       lastChild = lastChild.getTreePrev();
978     }
979     return lastChild;
980   }
981
982   @Override
983   public boolean isIncomplete() {
984     // if there's something following us, we're not incomplete
985     if (!PsiTreeUtil.hasErrorElements(myNode.getPsi())) {
986       PsiElement element = myNode.getPsi().getNextSibling();
987       while (element instanceof PsiWhiteSpace) {
988         element = element.getNextSibling();
989       }
990       if (element != null) {
991         return false;
992       }
993     }
994
995     final ASTNode lastChild = getLastNonSpaceChild(myNode, false);
996     if (lastChild != null) {
997       if (lastChild.getElementType() == PyElementTypes.STATEMENT_LIST) {
998         // only multiline statement lists are considered incomplete
999         final ASTNode statementListPrev = lastChild.getTreePrev();
1000         if (statementListPrev != null && statementListPrev.getText().indexOf('\n') >= 0) {
1001           return true;
1002         }
1003       }
1004       if (lastChild.getElementType() == PyElementTypes.BINARY_EXPRESSION) {
1005         final PyBinaryExpression binaryExpression = (PyBinaryExpression)lastChild.getPsi();
1006         if (binaryExpression.getRightExpression() == null) {
1007           return true;
1008         }
1009       }
1010       if (isIncompleteCall(lastChild)) return true;
1011     }
1012
1013     if (myNode.getPsi() instanceof PyArgumentList) {
1014       final PyArgumentList argumentList = (PyArgumentList)myNode.getPsi();
1015       return argumentList.getClosingParen() == null;
1016     }
1017     if (isIncompleteCall(myNode)) {
1018       return true;
1019     }
1020
1021     return false;
1022   }
1023
1024   private static boolean isIncompleteCall(@NotNull ASTNode node) {
1025     if (node.getElementType() == PyElementTypes.CALL_EXPRESSION) {
1026       final PyCallExpression callExpression = (PyCallExpression)node.getPsi();
1027       final PyArgumentList argumentList = callExpression.getArgumentList();
1028       if (argumentList == null || argumentList.getClosingParen() == null) {
1029         return true;
1030       }
1031     }
1032     return false;
1033   }
1034
1035   @Override
1036   public boolean isLeaf() {
1037     return myNode.getFirstChildNode() == null;
1038   }
1039 }