1cf3c42d8d6910044e36f79b5e4cc9e3882786ae
[idea/community.git] / 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 }