python/src/com/jetbrains/python/PythonFoldingBuilder.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;
  17
  18 import com.intellij.codeInsight.folding.CodeFoldingSettings;
  19 import com.intellij.lang.ASTNode;
  20 import com.intellij.lang.folding.CustomFoldingBuilder;
  21 import com.intellij.lang.folding.FoldingDescriptor;
  22 import com.intellij.openapi.editor.Document;
  23 import com.intellij.openapi.project.DumbAware;
  24 import com.intellij.openapi.util.Pair;
  25 import com.intellij.openapi.util.TextRange;
  26 import com.intellij.openapi.util.text.LineTokenizer;
  27 import com.intellij.openapi.util.text.StringUtil;
  28 import com.intellij.psi.PsiElement;
  29 import com.intellij.psi.PsiWhiteSpace;
  30 import com.intellij.psi.tree.IElementType;
  31 import com.intellij.psi.tree.TokenSet;
  32 import com.jetbrains.python.psi.*;
  33 import com.jetbrains.python.psi.impl.PyFileImpl;
  34 import org.jetbrains.annotations.NotNull;
  35 import org.jetbrains.annotations.Nullable;
  36
  37 import java.util.List;
  38
  39 /**
  40  * @author yole
  41  */
  42 public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAware {
  43
  44   public static final TokenSet FOLDABLE_COLLECTIONS_LITERALS = TokenSet.create(
  45                                                      PyElementTypes.SET_LITERAL_EXPRESSION,
  46                                                      PyElementTypes.DICT_LITERAL_EXPRESSION,
  47                                                      PyElementTypes.GENERATOR_EXPRESSION,
  48                                                      PyElementTypes.SET_COMP_EXPRESSION,
  49                                                      PyElementTypes.DICT_COMP_EXPRESSION,
  50                                                      PyElementTypes.LIST_LITERAL_EXPRESSION,
  51                                                      PyElementTypes.LIST_COMP_EXPRESSION,
  52                                                      PyElementTypes.TUPLE_EXPRESSION);
  53
  54   @Override
  55   protected void buildLanguageFoldRegions(@NotNull List<FoldingDescriptor> descriptors,
  56                                           @NotNull PsiElement root,
  57                                           @NotNull Document document,
  58                                           boolean quick) {
  59     appendDescriptors(root.getNode(), descriptors);
  60   }
  61
  62   private static void appendDescriptors(ASTNode node, List<FoldingDescriptor> descriptors) {
  63     IElementType elementType = node.getElementType();
  64     if (elementType instanceof PyFileElementType) {
  65       final List<PyImportStatementBase> imports = ((PyFile)node.getPsi()).getImportBlock();
  66       if (imports.size() > 1) {
  67         final PyImportStatementBase firstImport = imports.get(0);
  68         final PyImportStatementBase lastImport = imports.get(imports.size()-1);
  69         descriptors.add(new FoldingDescriptor(firstImport, new TextRange(firstImport.getTextRange().getStartOffset(),
  70                                                                          lastImport.getTextRange().getEndOffset())));
  71       }
  72     }
  73     else if (elementType == PyElementTypes.STATEMENT_LIST) {
  74       foldStatementList(node, descriptors);
  75     }
  76     else if (elementType == PyElementTypes.STRING_LITERAL_EXPRESSION) {
  77       foldLongStrings(node, descriptors);
  78     }
  79     else if (FOLDABLE_COLLECTIONS_LITERALS.contains(elementType)) {
  80       foldCollectionLiteral(node, descriptors);
  81     }
  82     else if (elementType == PyTokenTypes.END_OF_LINE_COMMENT) {
  83       foldSequentialComments(node, descriptors);
  84     }
  85     ASTNode child = node.getFirstChildNode();
  86     while (child != null) {
  87       appendDescriptors(child, descriptors);
  88       child = child.getTreeNext();
  89     }
  90   }
  91
  92   private static void foldSequentialComments(ASTNode node, List<FoldingDescriptor> descriptors) {
  93     //need to skip previous comments in sequence
  94     ASTNode curNode = node.getTreePrev();
  95     while (curNode != null) {
  96       if (curNode.getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
  97         return;
  98       }
  99       curNode = curNode.getPsi() instanceof PsiWhiteSpace ? curNode.getTreePrev() : null;
 100     }
 101
 102     //fold sequence comments in one block
 103     curNode = node.getTreeNext();
 104     ASTNode lastCommentNode = node;
 105     while (curNode != null) {
 106       if (curNode.getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
 107         lastCommentNode = curNode;
 108         curNode = curNode.getTreeNext();
 109         continue;
 110       }
 111       curNode = curNode.getPsi() instanceof PsiWhiteSpace ? curNode.getTreeNext() : null;
 112     }
 113
 114     if (lastCommentNode != node) {
 115       descriptors.add(new FoldingDescriptor(node, TextRange.create(node.getStartOffset(), lastCommentNode.getTextRange().getEndOffset())));
 116     }
 117
 118   }
 119
 120   private static void foldCollectionLiteral(ASTNode node, List<FoldingDescriptor> descriptors) {
 121     if (StringUtil.countNewLines(node.getChars()) > 0) {
 122       TextRange range = node.getTextRange();
 123       int delta = node.getElementType() == PyElementTypes.TUPLE_EXPRESSION ? 0 : 1;
 124       descriptors.add(new FoldingDescriptor(node, TextRange.create(range.getStartOffset() + delta, range.getEndOffset() - delta)));
 125     }
 126   }
 127
 128   private static void foldStatementList(ASTNode node, List<FoldingDescriptor> descriptors) {
 129     IElementType elType = node.getTreeParent().getElementType();
 130     if (elType == PyElementTypes.FUNCTION_DECLARATION
 131         || elType == PyElementTypes.CLASS_DECLARATION
 132         || ifFoldBlocks(node, elType)) {
 133       ASTNode colon = node.getTreeParent().findChildByType(PyTokenTypes.COLON);
 134       if (colon != null && colon.getStartOffset() + 1 < node.getTextRange().getEndOffset() - 1) {
 135         final CharSequence chars = node.getChars();
 136         int nodeStart = node.getTextRange().getStartOffset();
 137         int endOffset = node.getTextRange().getEndOffset();
 138         while(endOffset > colon.getStartOffset()+2 && endOffset > nodeStart && Character.isWhitespace(chars.charAt(endOffset - nodeStart - 1))) {
 139           endOffset--;
 140         }
 141         descriptors.add(new FoldingDescriptor(node, new TextRange(colon.getStartOffset() + 1, endOffset)));
 142       }
 143       else {
 144         TextRange range = node.getTextRange();
 145         if (range.getStartOffset() < range.getEndOffset() - 1) { // only for ranges at least 1 char wide
 146           descriptors.add(new FoldingDescriptor(node, range));
 147         }
 148       }
 149     }
 150   }
 151
 152   private static boolean ifFoldBlocks(ASTNode statementList, IElementType parentType) {
 153     if (!PyElementTypes.PARTS.contains(parentType)) {
 154       return false;
 155     }
 156     PsiElement element = statementList.getPsi();
 157     if (element instanceof PyStatementList) {
 158       return StringUtil.countNewLines(element.getText()) > 0;
 159     }
 160     return false;
 161   }
 162
 163   private static void foldLongStrings(ASTNode node, List<FoldingDescriptor> descriptors) {
 164     //don't want to fold docstrings like """\n string \n """
 165     boolean shouldFoldDocString = getDocStringOwnerType(node) != null && StringUtil.countNewLines(node.getChars()) > 1;
 166     boolean shouldFoldString = getDocStringOwnerType(node) == null && StringUtil.countNewLines(node.getChars()) > 0;
 167     if (shouldFoldDocString || shouldFoldString) {
 168       descriptors.add(new FoldingDescriptor(node, node.getTextRange()));
 169     }
 170   }
 171
 172   @Nullable
 173   private static IElementType getDocStringOwnerType(ASTNode node) {
 174     final ASTNode treeParent = node.getTreeParent();
 175     IElementType parentType = treeParent.getElementType();
 176     if (parentType == PyElementTypes.EXPRESSION_STATEMENT && treeParent.getTreeParent() != null) {
 177       final ASTNode parent2 = treeParent.getTreeParent();
 178       if (parent2.getElementType() == PyElementTypes.STATEMENT_LIST && parent2.getTreeParent() != null && treeParent == parent2.getFirstChildNode()) {
 179         final ASTNode parent3 = parent2.getTreeParent();
 180         if (parent3.getElementType() == PyElementTypes.FUNCTION_DECLARATION || parent3.getElementType() == PyElementTypes.CLASS_DECLARATION) {
 181           return parent3.getElementType();
 182         }
 183       }
 184       else if (parent2.getElementType() instanceof PyFileElementType) {
 185         return parent2.getElementType();
 186       }
 187     }
 188     return null;
 189   }
 190
 191   @Override
 192   protected String getLanguagePlaceholderText(@NotNull ASTNode node, @NotNull TextRange range) {
 193     if (PyFileImpl.isImport(node, false)) {
 194       return "import ...";
 195     }
 196     if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) {
 197       PyStringLiteralExpression stringLiteralExpression = (PyStringLiteralExpression)node.getPsi();
 198       if (stringLiteralExpression.isDocString()) {
 199         final String stringValue = stringLiteralExpression.getStringValue().trim();
 200         final String[] lines = LineTokenizer.tokenize(stringValue, true);
 201         if (lines.length > 2 && lines[1].trim().length() == 0) {
 202           return "\"\"\"" + lines[0].trim() + "...\"\"\"";
 203         }
 204         return "\"\"\"...\"\"\"";
 205       } else {
 206         return getLanguagePlaceholderForString(stringLiteralExpression);
 207       }
 208     }
 209     return "...";
 210   }
 211
 212   private static String getLanguagePlaceholderForString(PyStringLiteralExpression stringLiteralExpression) {
 213     String stringText = stringLiteralExpression.getText();
 214     Pair<String, String> quotes = PythonStringUtil.getQuotes(stringText);
 215     if (quotes != null) {
 216       return quotes.second + "..." + quotes.second;
 217     }
 218     return "...";
 219   }
 220
 221   @Override
 222   protected boolean isRegionCollapsedByDefault(@NotNull ASTNode node) {
 223     if (PyFileImpl.isImport(node, false)) {
 224       return CodeFoldingSettings.getInstance().COLLAPSE_IMPORTS;
 225     }
 226     if (node.getElementType() == PyElementTypes.STRING_LITERAL_EXPRESSION) {
 227       if (getDocStringOwnerType(node) == PyElementTypes.FUNCTION_DECLARATION && CodeFoldingSettings.getInstance().COLLAPSE_METHODS) {
 228         // method will be collapsed, no need to also collapse docstring
 229         return false;
 230       }
 231       if (getDocStringOwnerType(node) != null) {
 232         return CodeFoldingSettings.getInstance().COLLAPSE_DOC_COMMENTS;
 233       }
 234       return PythonFoldingSettings.getInstance().isCollapseLongStrings();
 235     }
 236     if (node.getElementType() == PyTokenTypes.END_OF_LINE_COMMENT) {
 237       return PythonFoldingSettings.getInstance().isCollapseSequentialComments();
 238     }
 239     if (node.getElementType() == PyElementTypes.STATEMENT_LIST && node.getTreeParent().getElementType() == PyElementTypes.FUNCTION_DECLARATION) {
 240       return CodeFoldingSettings.getInstance().COLLAPSE_METHODS;
 241     }
 242     if (FOLDABLE_COLLECTIONS_LITERALS.contains(node.getElementType())) {
 243       return PythonFoldingSettings.getInstance().isCollapseLongCollections();
 244     }
 245     return false;
 246   }
 247
 248   @Override
 249   protected boolean isCustomFoldingCandidate(ASTNode node) {
 250     return node.getElementType() == PyTokenTypes.END_OF_LINE_COMMENT;
 251   }
 252
 253   @Override
 254   protected boolean isCustomFoldingRoot(ASTNode node) {
 255     return node.getPsi() instanceof PyFile || node.getElementType() == PyElementTypes.STATEMENT_LIST;
 256   }
 257 }