All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.jetbrains.python.formatter.PyBlock Maven / Gradle / Ivy

Go to download

A packaging of the IntelliJ Community Edition python-community library. This is release number 1 of trunk branch 142.

The newest version!
/*
 * Copyright 2000-2014 JetBrains s.r.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.jetbrains.python.formatter;

import com.intellij.formatting.*;
import com.intellij.lang.ASTNode;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.util.Key;
import com.intellij.openapi.util.TextRange;
import com.intellij.psi.*;
import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
import com.intellij.psi.impl.source.tree.TreeUtil;
import com.intellij.psi.tree.IElementType;
import com.intellij.psi.tree.TokenSet;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.ArrayUtil;
import com.jetbrains.python.PyElementTypes;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.PythonDialectsTokenSetProvider;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_COLON;
import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_VALUE;
import static com.jetbrains.python.formatter.PythonFormattingModelBuilder.STATEMENT_OR_DECLARATION;

/**
 * @author yole
 */
public class PyBlock implements ASTBlock {
  private static final TokenSet ourListElementTypes = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION,
                                                                      PyElementTypes.LIST_COMP_EXPRESSION,
                                                                      PyElementTypes.DICT_COMP_EXPRESSION,
                                                                      PyElementTypes.SET_COMP_EXPRESSION,
                                                                      PyElementTypes.DICT_LITERAL_EXPRESSION,
                                                                      PyElementTypes.SET_LITERAL_EXPRESSION,
                                                                      PyElementTypes.ARGUMENT_LIST,
                                                                      PyElementTypes.PARAMETER_LIST,
                                                                      PyElementTypes.TUPLE_EXPRESSION,
                                                                      PyElementTypes.PARENTHESIZED_EXPRESSION,
                                                                      PyElementTypes.SLICE_EXPRESSION,
                                                                      PyElementTypes.SUBSCRIPTION_EXPRESSION,
                                                                      PyElementTypes.GENERATOR_EXPRESSION);

  private static final TokenSet ourBrackets = TokenSet.create(PyTokenTypes.LPAR, PyTokenTypes.RPAR,
                                                              PyTokenTypes.LBRACE, PyTokenTypes.RBRACE,
                                                              PyTokenTypes.LBRACKET, PyTokenTypes.RBRACKET);

  private static final TokenSet ourHangingIndentOwners = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION,
                                                                         PyElementTypes.DICT_LITERAL_EXPRESSION,
                                                                         PyElementTypes.SET_LITERAL_EXPRESSION,
                                                                         PyElementTypes.ARGUMENT_LIST,
                                                                         PyElementTypes.PARAMETER_LIST,
                                                                         PyElementTypes.TUPLE_EXPRESSION,
                                                                         PyElementTypes.PARENTHESIZED_EXPRESSION,
                                                                         PyElementTypes.GENERATOR_EXPRESSION,
                                                                         PyElementTypes.FUNCTION_DECLARATION,
                                                                         PyElementTypes.CALL_EXPRESSION,
                                                                         PyElementTypes.FROM_IMPORT_STATEMENT);

  private static final boolean DUMP_FORMATTING_BLOCKS = false;
  public static final Key IMPORT_GROUP_BEGIN = Key.create("com.jetbrains.python.formatter.importGroupBegin");

  private final PyBlock myParent;
  private final Alignment myAlignment;
  private final Indent myIndent;
  private final ASTNode myNode;
  private final Wrap myWrap;
  private final PyBlockContext myContext;
  private List mySubBlocks = null;
  private Map mySubBlockByNode = null;
  private Alignment myChildAlignment;
  private final Alignment myDictAlignment;
  private final Wrap myDictWrapping;
  private final boolean myEmptySequence;

  public PyBlock(final PyBlock parent,
                 final ASTNode node,
                 final Alignment alignment,
                 final Indent indent,
                 final Wrap wrap,
                 final PyBlockContext context) {
    myParent = parent;
    myAlignment = alignment;
    myIndent = indent;
    myNode = node;
    myWrap = wrap;
    myContext = context;
    myEmptySequence = isEmptySequence(node);

    if (node.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION) {
      myDictAlignment = Alignment.createAlignment(true);
      myDictWrapping = Wrap.createWrap(myContext.getPySettings().DICT_WRAPPING, true);
    }
    else {
      myDictAlignment = null;
      myDictWrapping = null;
    }
  }

  @NotNull
  public ASTNode getNode() {
    return myNode;
  }

  @NotNull
  public TextRange getTextRange() {
    return myNode.getTextRange();
  }

  private Alignment getAlignmentForChildren() {
    if (myChildAlignment == null) {
      myChildAlignment = Alignment.createAlignment();
    }
    return myChildAlignment;
  }

  @NotNull
  public List getSubBlocks() {
    if (mySubBlocks == null) {
      mySubBlockByNode = buildSubBlocks();
      mySubBlocks = new ArrayList(mySubBlockByNode.values());
      if (DUMP_FORMATTING_BLOCKS) {
        dumpSubBlocks();
      }
    }
    return Collections.unmodifiableList(mySubBlocks);
  }

  @Nullable
  private PyBlock getSubBlockByNode(@NotNull ASTNode node) {
    return mySubBlockByNode.get(node);
  }

  @Nullable
  private PyBlock getSubBlockByIndex(int index) {
    return mySubBlocks.get(index);
  }

  @NotNull
  private Map buildSubBlocks() {
    final Map blocks = new LinkedHashMap();
    for (ASTNode child = myNode.getFirstChildNode(); child != null; child = child.getTreeNext()) {

      final IElementType childType = child.getElementType();

      if (child.getTextRange().isEmpty()) continue;

      if (childType == TokenType.WHITE_SPACE) {
        continue;
      }

      blocks.put(child, buildSubBlock(child));
    }
    return Collections.unmodifiableMap(blocks);
  }

  private PyBlock buildSubBlock(ASTNode child) {
    final IElementType parentType = myNode.getElementType();

    final ASTNode grandParentNode = myNode.getTreeParent();
    final IElementType grandparentType = grandParentNode == null ? null : grandParentNode.getElementType();

    final IElementType childType = child.getElementType();
    Wrap wrap = null;
    Indent childIndent = Indent.getNoneIndent();
    Alignment childAlignment = null;

    if (parentType == PyElementTypes.BINARY_EXPRESSION && !isInControlStatement()) {
      //Setup alignments for binary expression
      childAlignment = getAlignmentForChildren();

      PyBlock p = myParent; //Check grandparents
      while (p != null) {
        final ASTNode pNode = p.getNode();
        if (ourListElementTypes.contains(pNode.getElementType())) {
          if (needListAlignment(child) && !myEmptySequence) {

            childAlignment = p.getChildAlignment();
            break;
          }
        }
        else if (pNode == PyElementTypes.BINARY_EXPRESSION) {
          childAlignment = p.getChildAlignment();
        }
        if (!breaksAlignment(pNode.getElementType())) {
          p = p.myParent;
        }
        else {
          break;
        }
      }
    }

    if (childType == PyElementTypes.STATEMENT_LIST) {
      if (hasLineBreaksBefore(child, 1) || needLineBreakInStatement()) {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (childType == PyElementTypes.IMPORT_ELEMENT) {
      wrap = Wrap.createWrap(WrapType.NORMAL, true);
      childIndent = Indent.getNormalIndent();
    }
    if (childType == PyTokenTypes.END_OF_LINE_COMMENT && parentType == PyElementTypes.FROM_IMPORT_STATEMENT) {
      childIndent = Indent.getNormalIndent();
    }
    if (ourListElementTypes.contains(parentType)) {
      // wrapping in non-parenthesized tuple expression is not allowed (PY-1792)
      if ((parentType != PyElementTypes.TUPLE_EXPRESSION || grandparentType == PyElementTypes.PARENTHESIZED_EXPRESSION) &&
          !ourBrackets.contains(childType) &&
          childType != PyTokenTypes.COMMA &&
          !isSliceOperand(child) /*&& !isSubscriptionOperand(child)*/) {
        wrap = Wrap.createWrap(WrapType.NORMAL, true);
      }
      if (needListAlignment(child) && !myEmptySequence) {
        childAlignment = getAlignmentForChildren();
      }
      if (childType == PyTokenTypes.END_OF_LINE_COMMENT) {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.BINARY_EXPRESSION &&
             (PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens().contains(childType) ||
              PyTokenTypes.OPERATIONS.contains(childType))) {
      if (isInControlStatement()) {
        final PyParenthesizedExpression parens = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyParenthesizedExpression.class, true,
                                                                       PyStatementPart.class);
        childIndent = parens != null ? Indent.getNormalIndent() : Indent.getContinuationIndent();
      }
    }

    PyCodeStyleSettings settings = CodeStyleSettingsManager.getSettings(child.getPsi().getProject()).getCustomSettings(PyCodeStyleSettings.class);
    if (parentType == PyElementTypes.LIST_LITERAL_EXPRESSION || parentType == PyElementTypes.LIST_COMP_EXPRESSION) {
      if (childType == PyTokenTypes.RBRACKET || childType == PyTokenTypes.LBRACKET) {
        childIndent = Indent.getNoneIndent();
      }
      else {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.DICT_LITERAL_EXPRESSION || parentType == PyElementTypes.SET_LITERAL_EXPRESSION ||
             parentType == PyElementTypes.SET_COMP_EXPRESSION || parentType == PyElementTypes.DICT_COMP_EXPRESSION) {
      if (childType == PyTokenTypes.RBRACE || !hasLineBreaksBefore(child, 1)) {
        childIndent = Indent.getNoneIndent();
      }
      else {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.STRING_LITERAL_EXPRESSION) {
      if (PyTokenTypes.STRING_NODES.contains(childType)) {
        childAlignment = getAlignmentForChildren();
      }
    }
    else if (parentType == PyElementTypes.FROM_IMPORT_STATEMENT) {
      if (myNode.findChildByType(PyTokenTypes.LPAR) != null) {
        if (childType == PyElementTypes.IMPORT_ELEMENT) {
          if (settings.ALIGN_MULTILINE_IMPORTS) {
            childAlignment = getAlignmentForChildren();
          }
          else {
            childIndent = Indent.getNormalIndent();
          }
        }
        if (childType == PyTokenTypes.RPAR) {
          childIndent = Indent.getNoneIndent();
          if (!hasHangingIndent(myNode.getPsi())) {
            childAlignment = getAlignmentForChildren();
          }
        }
      }
    }
    else if (isValueOfKeyValuePair(child)) {
      childIndent = Indent.getNormalIndent();
    }
    //Align elements vertically if there is an argument in the first line of parenthesized expression
    else if (!hasHangingIndent(myNode.getPsi()) &&
             ((parentType == PyElementTypes.PARENTHESIZED_EXPRESSION && myContext.getSettings().ALIGN_MULTILINE_PARENTHESIZED_EXPRESSION) ||
              (parentType == PyElementTypes.ARGUMENT_LIST && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS) ||
              (parentType == PyElementTypes.PARAMETER_LIST && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS)) &&
             !isIndentNext(child) &&
             !hasLineBreaksBefore(myNode.getFirstChildNode(), 1) &&
             !ourListElementTypes.contains(childType)) {

      if (!ourBrackets.contains(childType)) {
        childAlignment = getAlignmentForChildren();
        if (parentType != PyElementTypes.CALL_EXPRESSION) {
          childIndent = Indent.getNormalIndent();
        }
      }
      else if (childType == PyTokenTypes.RPAR) {
        childIndent = Indent.getNoneIndent();
      }
    }
    else if (parentType == PyElementTypes.GENERATOR_EXPRESSION || parentType == PyElementTypes.PARENTHESIZED_EXPRESSION) {
      if (childType == PyTokenTypes.RPAR || !hasLineBreaksBefore(child, 1)) {
        childIndent = Indent.getNoneIndent();
      }
      else {
        childIndent = isIndentNext(child) ? Indent.getContinuationIndent() : Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.ARGUMENT_LIST || parentType == PyElementTypes.PARAMETER_LIST) {
      if (childType == PyTokenTypes.RPAR) {
        childIndent = Indent.getNoneIndent();
      }
      else {
        if (parentType == PyElementTypes.PARAMETER_LIST || argumentMayHaveSameIndentAsFollowingStatementList()) {
          childIndent = Indent.getContinuationIndent();
        }
        else {
          childIndent = Indent.getNormalIndent();
        }
      }
    }
    else if (parentType == PyElementTypes.SUBSCRIPTION_EXPRESSION) {
      final PyExpression indexExpression = ((PySubscriptionExpression)myNode.getPsi()).getIndexExpression();
      if (indexExpression != null && child == indexExpression.getNode()) {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.REFERENCE_EXPRESSION) {
      if (child != myNode.getFirstChildNode()) {
        childIndent = Indent.getNormalIndent();
        if (hasLineBreaksBefore(child, 1)) {
          if (isInControlStatement()) {
            childIndent = Indent.getContinuationIndent();
          }
          else {
            PyBlock b = myParent;
            while (b != null) {
              if (b.getNode().getPsi() instanceof PyParenthesizedExpression ||
                  b.getNode().getPsi() instanceof PyArgumentList ||
                  b.getNode().getPsi() instanceof PyParameterList) {
                childAlignment = getAlignmentOfChild(b, 1);
                break;
              }
              b = b.myParent;
            }
          }
        }
      }
    }
    if (childType == PyElementTypes.KEY_VALUE_EXPRESSION && isChildOfDictLiteral(child)) {
      wrap = myDictWrapping;
      childIndent = Indent.getNormalIndent();
    }

    if (isAfterStatementList(child) && !hasLineBreaksBefore(child, 2) && child.getElementType() != PyTokenTypes.END_OF_LINE_COMMENT) {
      // maybe enter was pressed and cut us from a previous (nested) statement list
      childIndent = Indent.getNormalIndent();
    }

    if (settings.DICT_ALIGNMENT == DICT_ALIGNMENT_ON_VALUE) {
      if (isValueOfKeyValuePairOfDictLiteral(child) && !ourListElementTypes.contains(childType)) {
        childAlignment = myParent.myDictAlignment;
      }
      else if (isValueOfKeyValuePairOfDictLiteral(myNode) &&
               ourListElementTypes.contains(parentType) &&
               PyTokenTypes.OPEN_BRACES.contains(childType)) {
        childAlignment = myParent.myParent.myDictAlignment;
      }
    }
    else if (myContext.getPySettings().DICT_ALIGNMENT == DICT_ALIGNMENT_ON_COLON) {
      if (isChildOfKeyValuePairOfDictLiteral(child) && childType == PyTokenTypes.COLON) {
        childAlignment = myParent.myDictAlignment;
      }
    }

    ASTNode prev = child.getTreePrev();
    while (prev != null && prev.getElementType() == TokenType.WHITE_SPACE) {
      if (prev.textContains('\\') &&
          !childIndent.equals(Indent.getContinuationIndent(false)) &&
          !childIndent.equals(Indent.getContinuationIndent(true))) {
        childIndent = isIndentNext(child) ? Indent.getContinuationIndent() : Indent.getNormalIndent();
        break;
      }
      prev = prev.getTreePrev();
    }

    return new PyBlock(this, child, childAlignment, childIndent, wrap, myContext);
  }

  private static boolean isValueOfKeyValuePairOfDictLiteral(@NotNull ASTNode node) {
    return isValueOfKeyValuePair(node) && isChildOfDictLiteral(node.getTreeParent());
  }

  private static boolean isChildOfKeyValuePairOfDictLiteral(@NotNull ASTNode node) {
    return isChildOfKeyValuePair(node) && isChildOfDictLiteral(node.getTreeParent());
  }

  private static boolean isChildOfDictLiteral(@NotNull ASTNode node) {
    final ASTNode nodeParent = node.getTreeParent();
    return nodeParent != null && nodeParent.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION;
  }

  private static boolean isChildOfKeyValuePair(@NotNull ASTNode node) {
    final ASTNode nodeParent = node.getTreeParent();
    return nodeParent != null && nodeParent.getElementType() == PyElementTypes.KEY_VALUE_EXPRESSION;
  }

  private static boolean isValueOfKeyValuePair(@NotNull ASTNode node) {
    return isChildOfKeyValuePair(node) && node.getTreeParent().getPsi(PyKeyValueExpression.class).getValue() == node.getPsi();
  }

  private static boolean isEmptySequence(@NotNull ASTNode node) {
    return node.getPsi() instanceof PySequenceExpression && ((PySequenceExpression)node.getPsi()).isEmpty();
  }

  private boolean argumentMayHaveSameIndentAsFollowingStatementList() {
    // This check is supposed to prevent PEP8's error: Continuation line with the same indent as next logical line
    final PsiElement header = getControlStatementHeader(myNode);
    if (header instanceof PyStatementListContainer) {
      final PyStatementList statementList = ((PyStatementListContainer)header).getStatementList();
      final int headerStartLine = getLineInDocument(header);
      final int statementListStartLine = getLineInDocument(statementList);
      final int argumentListStartLine = getLineInDocument(myNode.getPsi());
      return headerStartLine == argumentListStartLine && headerStartLine != statementListStartLine;
    }
    return false;
  }

  // Check https://www.python.org/dev/peps/pep-0008/#indentation
  private static boolean hasHangingIndent(@NotNull PsiElement elem) {
    if (elem instanceof PyCallExpression) {
      final PyArgumentList argumentList = ((PyCallExpression)elem).getArgumentList();
      return argumentList != null && hasHangingIndent(argumentList);
    }
    else if (elem instanceof PyFunction) {
      return hasHangingIndent(((PyFunction)elem).getParameterList());
    }

    final PsiElement firstChild;
    if (elem instanceof PyFromImportStatement) {
      firstChild = ((PyFromImportStatement)elem).getLeftParen();
    }
    else {
      firstChild = elem.getFirstChild();
    }
    if (firstChild == null) {
      return false;
    }
    final IElementType elementType = elem.getNode().getElementType();
    final ASTNode firstChildNode = firstChild.getNode();
    if (ourHangingIndentOwners.contains(elementType) && PyTokenTypes.OPEN_BRACES.contains(firstChildNode.getElementType())) {
      if (hasLineBreaksAfter(firstChildNode, 1)) {
        return true;
      }
      final PsiElement[] items = getItems(elem);
      if (items.length == 0) {
        return !PyTokenTypes.CLOSE_BRACES.contains(elem.getLastChild().getNode().getElementType());
      }
      else {
        final PsiElement firstItem = items[0];
        if (firstItem instanceof PyNamedParameter) {
          final PyExpression defaultValue = ((PyNamedParameter)firstItem).getDefaultValue();
          return defaultValue != null && hasHangingIndent(defaultValue);
        }
        else if (firstItem instanceof PyKeywordArgument) {
          final PyExpression valueExpression = ((PyKeywordArgument)firstItem).getValueExpression();
          return valueExpression != null && hasHangingIndent(valueExpression);
        }
        else if (firstItem instanceof PyKeyValueExpression) {
          final PyExpression value = ((PyKeyValueExpression)firstItem).getValue();
          return value != null && hasHangingIndent(value);
        }
        return hasHangingIndent(firstItem);
      }
    }
    else {
      return false;
    }
  }

  @NotNull
  private static PsiElement[] getItems(@NotNull PsiElement elem) {
    if (elem instanceof PySequenceExpression) {
      return ((PySequenceExpression)elem).getElements();
    }
    else if (elem instanceof PyParameterList) {
      return ((PyParameterList)elem).getParameters();
    }
    else if (elem instanceof PyArgumentList) {
      return ((PyArgumentList)elem).getArguments();
    }
    else if (elem instanceof PyFromImportStatement) {
      return ((PyFromImportStatement)elem).getImportElements();
    }
    else if (elem instanceof PyParenthesizedExpression) {
      final PyExpression containedExpression = ((PyParenthesizedExpression)elem).getContainedExpression();
      if (containedExpression instanceof PyTupleExpression) {
        return ((PyTupleExpression)containedExpression).getElements();
      }
      else if (containedExpression != null) {
        return new PsiElement[]{containedExpression};
      }
    }
    return PsiElement.EMPTY_ARRAY;
  }

  private static boolean breaksAlignment(IElementType type) {
    return type != PyElementTypes.BINARY_EXPRESSION;
  }

  private static Alignment getAlignmentOfChild(PyBlock b, int childNum) {
    if (b.getSubBlocks().size() > childNum) {
      final ChildAttributes attributes = b.getChildAttributes(childNum);
      return attributes.getAlignment();
    }
    return null;
  }

  private static boolean isIndentNext(ASTNode child) {
    final PsiElement psi = PsiTreeUtil.getParentOfType(child.getPsi(), PyStatement.class);

    return psi instanceof PyIfStatement ||
           psi instanceof PyForStatement ||
           psi instanceof PyWithStatement ||
           psi instanceof PyClass ||
           psi instanceof PyFunction ||
           psi instanceof PyTryExceptStatement ||
           psi instanceof PyElsePart ||
           psi instanceof PyIfPart ||
           psi instanceof PyWhileStatement;
  }

  private static boolean isSubscriptionOperand(ASTNode child) {
    return child.getTreeParent().getElementType() == PyElementTypes.SUBSCRIPTION_EXPRESSION &&
           child.getPsi() == ((PySubscriptionExpression)child.getTreeParent().getPsi()).getOperand();
  }

  private boolean isInControlStatement() {
    return getControlStatementHeader(myNode) != null;
  }

  @Nullable
  private static PsiElement getControlStatementHeader(@NotNull ASTNode node) {
    final PyStatementPart statementPart = PsiTreeUtil.getParentOfType(node.getPsi(), PyStatementPart.class, false, PyStatementList.class);
    if (statementPart != null) {
      return statementPart;
    }
    final PyWithItem withItem = PsiTreeUtil.getParentOfType(node.getPsi(), PyWithItem.class);
    if (withItem != null) {
      return withItem.getParent();
    }
    return null;
  }

  private static int getLineInDocument(@NotNull PsiElement element) {
    final Document document = PsiDocumentManager.getInstance(element.getProject()).getDocument(element.getContainingFile());
    return document != null ? document.getLineNumber(element.getTextOffset()) : -1;
  }

  private boolean isSliceOperand(ASTNode child) {
    if (myNode.getPsi() instanceof PySliceExpression) {
      final PySliceExpression sliceExpression = (PySliceExpression)myNode.getPsi();
      final PyExpression operand = sliceExpression.getOperand();
      return operand.getNode() == child;
    }
    return false;
  }

  private static boolean isAfterStatementList(ASTNode child) {
    final PsiElement prev = child.getPsi().getPrevSibling();
    if (!(prev instanceof PyStatement)) {
      return false;
    }
    final PsiElement lastChild = PsiTreeUtil.getDeepestLast(prev);
    return lastChild.getParent() instanceof PyStatementList;
  }

  private boolean needListAlignment(ASTNode child) {
    final IElementType childType = child.getElementType();
    if (PyTokenTypes.OPEN_BRACES.contains(childType)) {
      return false;
    }
    if (PyTokenTypes.CLOSE_BRACES.contains(childType)) {
      final ASTNode prevNonSpace = findPrevNonSpaceNode(child);
      if (prevNonSpace != null &&
          prevNonSpace.getElementType() == PyTokenTypes.COMMA &&
          myContext.getMode() == FormattingMode.ADJUST_INDENT) {
        return true;
      }
      return !hasHangingIndent(myNode.getPsi());
    }
    if (myNode.getElementType() == PyElementTypes.ARGUMENT_LIST) {
      if (!myContext.getSettings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS || hasHangingIndent(myNode.getPsi())) {
        return false;
      }
      if (child.getElementType() == PyTokenTypes.COMMA) {
        return false;
      }
      return true;
    }
    if (myNode.getElementType() == PyElementTypes.PARAMETER_LIST) {
      return !hasHangingIndent(myNode.getPsi()) && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS;
    }
    if (myNode.getElementType() == PyElementTypes.SUBSCRIPTION_EXPRESSION) {
      return false;
    }
    if (child.getElementType() == PyTokenTypes.COMMA) {
      return false;
    }
    return myContext.getPySettings().ALIGN_COLLECTIONS_AND_COMPREHENSIONS && !hasHangingIndent(myNode.getPsi());
  }

  @Nullable
  private static ASTNode findPrevNonSpaceNode(ASTNode node) {
    do {
      node = node.getTreePrev();
    }
    while (isWhitespace(node));
    return node;
  }

  private static boolean isWhitespace(@Nullable ASTNode node) {
    return node != null && (node.getElementType() == TokenType.WHITE_SPACE || PyTokenTypes.WHITESPACE.contains(node.getElementType()));
  }

  private static boolean hasLineBreaksBefore(@NotNull ASTNode child, int minCount) {
    final ASTNode treePrev = child.getTreePrev();
    return (treePrev != null && isWhitespaceWithLineBreaks(TreeUtil.findLastLeaf(treePrev), minCount)) ||
           isWhitespaceWithLineBreaks(child.getFirstChildNode(), minCount);
  }

  private static boolean hasLineBreaksAfter(@NotNull ASTNode child, int minCount) {
    final ASTNode treeNext = child.getTreeNext();
    return (treeNext != null && isWhitespaceWithLineBreaks(TreeUtil.findLastLeaf(treeNext), minCount)) ||
           isWhitespaceWithLineBreaks(child.getLastChildNode(), minCount);
  }

  private static boolean isWhitespaceWithLineBreaks(@Nullable ASTNode node, int minCount) {
    if (isWhitespace(node)) {
      final String prevNodeText = node.getText();
      int count = 0;
      for (int i = 0; i < prevNodeText.length(); i++) {
        if (prevNodeText.charAt(i) == '\n') {
          count++;
          if (count == minCount) {
            return true;
          }
        }
      }
    }
    return false;
  }

  private void dumpSubBlocks() {
    System.out.println("Subblocks of " + myNode.getPsi() + ":");
    for (Block block : mySubBlocks) {
      if (block instanceof PyBlock) {
        System.out.println("  " + ((PyBlock)block).getNode().getPsi().toString() + " " + block.getTextRange().getStartOffset() + ":" + block
          .getTextRange().getLength());
      }
      else {
        System.out.println("  ");
      }
    }
  }

  @Nullable
  public Wrap getWrap() {
    return myWrap;
  }

  @Nullable
  public Indent getIndent() {
    assert myIndent != null;
    return myIndent;
  }

  @Nullable
  public Alignment getAlignment() {
    return myAlignment;
  }

  @Nullable
  public Spacing getSpacing(Block child1, @NotNull Block child2) {
    if (child1 instanceof ASTBlock && child2 instanceof ASTBlock) {
      final ASTNode node1 = ((ASTBlock)child1).getNode();
      ASTNode node2 = ((ASTBlock)child2).getNode();
      final IElementType childType1 = node1.getElementType();
      final PsiElement psi1 = node1.getPsi();

      PsiElement psi2 = node2.getPsi();
      // skip not inline comments to handles blank lines between various declarations
      if (psi2 instanceof PsiComment && hasLineBreaksBefore(node2, 1)) {
        final PsiElement nonCommentAfter = PyPsiUtils.getNextNonCommentSibling(psi2, true);
        if (nonCommentAfter != null) {
          psi2 = nonCommentAfter;
        }
      }
      node2 = psi2.getNode();
      final IElementType childType2 = psi2.getNode().getElementType();
      //noinspection ConstantConditions
      child2 = getSubBlockByNode(node2);

      final CommonCodeStyleSettings settings = myContext.getSettings();
      if (childType1 == PyTokenTypes.COLON && psi2 instanceof PyStatementList) {
        if (needLineBreakInStatement()) {
          return Spacing.createSpacing(0, 0, 1, true, settings.KEEP_BLANK_LINES_IN_CODE);
        }
      }

      if ((PyElementTypes.CLASS_OR_FUNCTION.contains(childType1) && STATEMENT_OR_DECLARATION.contains(childType2)) ||
          STATEMENT_OR_DECLARATION.contains(childType1) && PyElementTypes.CLASS_OR_FUNCTION.contains(childType2)) {
        if (PyUtil.isTopLevel(psi1)) {
          return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS);
        }
      }

      if (psi1 instanceof PyImportStatementBase) {
        if (psi2 instanceof PyImportStatementBase &&
            psi2.getCopyableUserData(IMPORT_GROUP_BEGIN) != null) {
          return Spacing.createSpacing(0, 0, 2, true, 1);
        }
        if (psi2 instanceof PyStatement && !(psi2 instanceof PyImportStatementBase)) {
          if (PyUtil.isTopLevel(psi1)) {
            return getBlankLinesForOption(settings.BLANK_LINES_AFTER_IMPORTS);
          }
          else {
            return getBlankLinesForOption(myContext.getPySettings().BLANK_LINES_AFTER_LOCAL_IMPORTS);
          }
        }
      }

      if (psi2 instanceof PsiComment && !hasLineBreaksBefore(psi2.getNode(), 1) && myContext.getPySettings().SPACE_BEFORE_NUMBER_SIGN) {
        return Spacing.createSpacing(2, 0, 0, false, 0);
      }
    }
    return myContext.getSpacingBuilder().getSpacing(this, child1, child2);
  }

  private Spacing getBlankLinesForOption(final int option) {
    final int blankLines = option + 1;
    return Spacing.createSpacing(0, 0, blankLines,
                                 myContext.getSettings().KEEP_LINE_BREAKS,
                                 myContext.getSettings().KEEP_BLANK_LINES_IN_DECLARATIONS);
  }

  private boolean needLineBreakInStatement() {
    final PyStatement statement = PsiTreeUtil.getParentOfType(myNode.getPsi(), PyStatement.class);
    if (statement != null) {
      final Collection parts = PsiTreeUtil.collectElementsOfType(statement, PyStatementPart.class);
      return (parts.size() == 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON) ||
             (parts.size() > 1 && myContext.getPySettings().NEW_LINE_AFTER_COLON_MULTI_CLAUSE);
    }
    return false;
  }

  @NotNull
  public ChildAttributes getChildAttributes(int newChildIndex) {
    int statementListsBelow = 0;
    if (newChildIndex > 0) {
      // always pass decision to a sane block from top level from file or definition
      if (myNode.getPsi() instanceof PyFile || myNode.getElementType() == PyTokenTypes.COLON) {
        return ChildAttributes.DELEGATE_TO_PREV_CHILD;
      }

      final PyBlock insertAfterBlock = getSubBlockByIndex(newChildIndex - 1);

      final ASTNode prevNode = insertAfterBlock.getNode();
      final PsiElement prevElt = prevNode.getPsi();

      // stmt lists, parts and definitions should also think for themselves
      if (prevElt instanceof PyStatementList) {
        if (dedentAfterLastStatement((PyStatementList)prevElt)) {
          return new ChildAttributes(Indent.getNoneIndent(), getChildAlignment());
        }
        return ChildAttributes.DELEGATE_TO_PREV_CHILD;
      }
      else if (prevElt instanceof PyStatementPart) {
        return ChildAttributes.DELEGATE_TO_PREV_CHILD;
      }

      ASTNode lastChild = insertAfterBlock.getNode();

      // HACK? This code fragment is needed to make testClass2() pass,
      // but I don't quite understand why it is necessary and why the formatter
      // doesn't request childAttributes from the correct block
      while (lastChild != null) {
        final IElementType last_type = lastChild.getElementType();
        if (last_type == PyElementTypes.STATEMENT_LIST && hasLineBreaksBefore(lastChild, 1)) {
          if (dedentAfterLastStatement((PyStatementList)lastChild.getPsi())) {
            break;
          }
          statementListsBelow++;
        }
        else if (statementListsBelow > 0 && lastChild.getPsi() instanceof PsiErrorElement) {
          statementListsBelow++;
        }
        if (myNode.getElementType() == PyElementTypes.STATEMENT_LIST && lastChild.getPsi() instanceof PsiErrorElement) {
          return ChildAttributes.DELEGATE_TO_PREV_CHILD;
        }
        lastChild = getLastNonSpaceChild(lastChild, true);
      }
    }

    // HACKETY-HACK
    // If a multi-step dedent follows the cursor position (see testMultiDedent()),
    // the whitespace (which must be a single Py:LINE_BREAK token) gets attached
    // to the outermost indented block (because we may not consume the DEDENT
    // tokens while parsing inner blocks). The policy is to put the indent to
    // the innermost block, so we need to resolve the situation here. Nested
    // delegation sometimes causes NPEs in formatter core, so we calculate the
    // correct indent manually.
    if (statementListsBelow > 0) { // was 1... strange
      @SuppressWarnings("ConstantConditions") final
      int indent = myContext.getSettings().getIndentOptions().INDENT_SIZE;
      return new ChildAttributes(Indent.getSpaceIndent(indent * statementListsBelow), null);
    }

    /*
    // it might be something like "def foo(): # comment" or "[1, # comment"; jump up to the real thing
    if (_node instanceof PsiComment || _node instanceof PsiWhiteSpace) {
      get
    }
    */


    final Indent childIndent = getChildIndent(newChildIndex);
    final Alignment childAlignment = getChildAlignment();
    return new ChildAttributes(childIndent, childAlignment);
  }

  private static boolean dedentAfterLastStatement(PyStatementList statementList) {
    final PyStatement[] statements = statementList.getStatements();
    if (statements.length == 0) {
      return false;
    }
    final PyStatement last = statements[statements.length - 1];
    return last instanceof PyReturnStatement || last instanceof PyRaiseStatement || last instanceof PyPassStatement;
  }

  @Nullable
  private Alignment getChildAlignment() {
    if (ourListElementTypes.contains(myNode.getElementType())) {
      if (isInControlStatement()) {
        return null;
      }
      if (myNode.getPsi() instanceof PyParameterList && !myContext.getSettings().ALIGN_MULTILINE_PARAMETERS) {
        return null;
      }
      if (myNode.getPsi() instanceof PyDictLiteralExpression) {
        final PyKeyValueExpression lastElement = ArrayUtil.getLastElement(((PyDictLiteralExpression)myNode.getPsi()).getElements());
        if (lastElement == null || lastElement.getValue() == null /* incomplete */) {
          return null;
        }
      }
      return getAlignmentForChildren();
    }
    return null;
  }

  private Indent getChildIndent(int newChildIndex) {
    final ASTNode afterNode = getAfterNode(newChildIndex);
    final ASTNode lastChild = getLastNonSpaceChild(myNode, false);
    if (lastChild != null && lastChild.getElementType() == PyElementTypes.STATEMENT_LIST && mySubBlocks.size() >= newChildIndex) {
      if (afterNode == null) {
        return Indent.getNoneIndent();
      }

      // handle pressing Enter after colon and before first statement in
      // existing statement list
      if (afterNode.getElementType() == PyElementTypes.STATEMENT_LIST || afterNode.getElementType() == PyTokenTypes.COLON) {
        return Indent.getNormalIndent();
      }

      // handle pressing Enter after colon when there is nothing in the
      // statement list
      final ASTNode lastFirstChild = lastChild.getFirstChildNode();
      if (lastFirstChild != null && lastFirstChild == lastChild.getLastChildNode() && lastFirstChild.getPsi() instanceof PsiErrorElement) {
        return Indent.getNormalIndent();
      }
    }
    else if (lastChild != null && PyElementTypes.LIST_LIKE_EXPRESSIONS.contains(lastChild.getElementType())) {
      // handle pressing enter at the end of a list literal when there's no closing paren or bracket 
      final ASTNode lastLastChild = lastChild.getLastChildNode();
      if (lastLastChild != null && lastLastChild.getPsi() instanceof PsiErrorElement) {
        // we're at a place like this: [foo, ... bar, 
        // we'd rather align to foo. this may be not a multiple of tabs.
        final PsiElement expr = lastChild.getPsi();
        PsiElement exprItem = expr.getFirstChild();
        boolean found = false;
        while (exprItem != null) { // find a worthy element to align to
          if (exprItem instanceof PyElement) {
            found = true; // align to foo in "[foo,"
            break;
          }
          if (exprItem instanceof PsiComment) {
            found = true; // align to foo in "[ # foo,"
            break;
          }
          exprItem = exprItem.getNextSibling();
        }
        if (found) {
          final PsiDocumentManager docMgr = PsiDocumentManager.getInstance(exprItem.getProject());
          final Document doc = docMgr.getDocument(exprItem.getContainingFile());
          if (doc != null) {
            int line_num = doc.getLineNumber(exprItem.getTextOffset());
            final int item_col = exprItem.getTextOffset() - doc.getLineStartOffset(line_num);
            final PsiElement here_elt = getNode().getPsi();
            line_num = doc.getLineNumber(here_elt.getTextOffset());
            final int node_col = here_elt.getTextOffset() - doc.getLineStartOffset(line_num);
            final int padding = item_col - node_col;
            if (padding > 0) { // negative is a syntax error,  but possible
              return Indent.getSpaceIndent(padding);
            }
          }
        }
        return Indent.getContinuationIndent(); // a fallback
      }
    }

    if (afterNode != null && afterNode.getElementType() == PyElementTypes.KEY_VALUE_EXPRESSION) {
      final PyKeyValueExpression keyValue = (PyKeyValueExpression)afterNode.getPsi();
      if (keyValue != null && keyValue.getValue() == null) {  // incomplete
        return Indent.getContinuationIndent();
      }
    }

    // constructs that imply indent for their children
    if (myNode.getElementType().equals(PyElementTypes.PARAMETER_LIST)) {
      return Indent.getContinuationIndent();
    }
    if (ourListElementTypes.contains(myNode.getElementType()) || myNode.getPsi() instanceof PyStatementPart) {
      return Indent.getNormalIndent();
    }

    if (afterNode != null) {
      ASTNode wsAfter = afterNode.getTreeNext();
      while (wsAfter != null && wsAfter.getElementType() == TokenType.WHITE_SPACE) {
        if (wsAfter.getText().indexOf('\\') >= 0) {
          return Indent.getNormalIndent();
        }
        wsAfter = wsAfter.getTreeNext();
      }
    }
    return Indent.getNoneIndent();
  }

  @Nullable
  private ASTNode getAfterNode(int newChildIndex) {
    if (newChildIndex == 0) {  // block text contains backslash line wrappings, child block list not built
      return null;
    }
    int prevIndex = newChildIndex - 1;
    while (prevIndex > 0 && getSubBlockByIndex(prevIndex).getNode().getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
      prevIndex--;
    }
    return getSubBlockByIndex(prevIndex).getNode();
  }

  private static ASTNode getLastNonSpaceChild(ASTNode node, boolean acceptError) {
    ASTNode lastChild = node.getLastChildNode();
    while (lastChild != null &&
           (lastChild.getElementType() == TokenType.WHITE_SPACE || (!acceptError && lastChild.getPsi() instanceof PsiErrorElement))) {
      lastChild = lastChild.getTreePrev();
    }
    return lastChild;
  }

  public boolean isIncomplete() {
    // if there's something following us, we're not incomplete
    if (!PsiTreeUtil.hasErrorElements(myNode.getPsi())) {
      PsiElement element = myNode.getPsi().getNextSibling();
      while (element instanceof PsiWhiteSpace) {
        element = element.getNextSibling();
      }
      if (element != null) {
        return false;
      }
    }

    final ASTNode lastChild = getLastNonSpaceChild(myNode, false);
    if (lastChild != null) {
      if (lastChild.getElementType() == PyElementTypes.STATEMENT_LIST) {
        // only multiline statement lists are considered incomplete
        final ASTNode statementListPrev = lastChild.getTreePrev();
        if (statementListPrev != null && statementListPrev.getText().indexOf('\n') >= 0) {
          return true;
        }
      }
      if (lastChild.getElementType() == PyElementTypes.BINARY_EXPRESSION) {
        final PyBinaryExpression binaryExpression = (PyBinaryExpression)lastChild.getPsi();
        if (binaryExpression.getRightExpression() == null) {
          return true;
        }
      }
      if (isIncompleteCall(lastChild)) return true;
    }

    if (myNode.getPsi() instanceof PyArgumentList) {
      final PyArgumentList argumentList = (PyArgumentList)myNode.getPsi();
      return argumentList.getClosingParen() == null;
    }
    if (isIncompleteCall(myNode)) {
      return true;
    }

    return false;
  }

  private static boolean isIncompleteCall(ASTNode node) {
    if (node.getElementType() == PyElementTypes.CALL_EXPRESSION) {
      final PyCallExpression callExpression = (PyCallExpression)node.getPsi();
      final PyArgumentList argumentList = callExpression.getArgumentList();
      if (argumentList == null || argumentList.getClosingParen() == null) {
        return true;
      }
    }
    return false;
  }

  public boolean isLeaf() {
    return myNode.getFirstChildNode() == null;
  }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy