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