Python custom folding enabled

python/src/com/jetbrains/python/PythonFoldingBuilder.java

@@ -2,6 +2,7 @@ package com.jetbrains.python;
 
 import com.intellij.codeInsight.folding.CodeFoldingSettings;
 import com.intellij.lang.ASTNode;
+import com.intellij.lang.folding.CustomFoldingBuilder;
 import com.intellij.lang.folding.FoldingBuilder;
 import com.intellij.lang.folding.FoldingDescriptor;
 import com.intellij.openapi.editor.Document;
@@ -9,8 +10,10 @@ import com.intellij.openapi.project.DumbAware;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.LineTokenizer;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.TokenType;
 import com.intellij.psi.tree.IElementType;
+import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.psi.PyFileElementType;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
 import org.jetbrains.annotations.NotNull;
@@ -22,12 +25,14 @@ import java.util.List;
 /**
  * @author yole
  */
-public class PythonFoldingBuilder implements FoldingBuilder, DumbAware {
-  @NotNull
-  public FoldingDescriptor[] buildFoldRegions(@NotNull ASTNode node, @NotNull Document document) {
-    List<FoldingDescriptor> descriptors = new ArrayList<FoldingDescriptor>();
-    appendDescriptors(node, descriptors);
-    return descriptors.toArray(new FoldingDescriptor[descriptors.size()]);
+public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAware {
+
+  @Override
+  protected void buildLanguageFoldRegions(@NotNull List<FoldingDescriptor> descriptors,
+                                          @NotNull PsiElement root,
+                                          @NotNull Document document,
+                                          boolean quick) {
+    appendDescriptors(root.getNode(), descriptors);
   }
 
   private static void appendDescriptors(ASTNode node, List<FoldingDescriptor> descriptors) {
@@ -122,7 +127,8 @@ public class PythonFoldingBuilder implements FoldingBuilder, DumbAware {
     return elementType == PyElementTypes.IMPORT_STATEMENT || elementType == PyElementTypes.FROM_IMPORT_STATEMENT;
   }
 
-  public String getPlaceholderText(@NotNull ASTNode node) {
+  @Override
+  protected String getLanguagePlaceholderText(@NotNull ASTNode node, @NotNull TextRange range) {
     if (isImport(node, false)) {
       return "import ...";
     }
@@ -137,7 +143,8 @@ public class PythonFoldingBuilder implements FoldingBuilder, DumbAware {
     return "...";
   }
 
-  public boolean isCollapsedByDefault(@NotNull ASTNode node) {
+  @Override
+  protected boolean isRegionCollapsedByDefault(@NotNull ASTNode node) {
     if (isImport(node, false)) {
       return CodeFoldingSettings.getInstance().COLLAPSE_IMPORTS;
     }
@@ -153,4 +160,14 @@ public class PythonFoldingBuilder implements FoldingBuilder, DumbAware {
     }
     return false;
   }
+
+  @Override
+  protected boolean isCustomFoldingCandidate(ASTNode node) {
+    return node.getElementType() == PyTokenTypes.END_OF_LINE_COMMENT;
+  }
+
+  @Override
+  protected boolean mayContainCustomFoldings(ASTNode node) {
+    return node.getPsi() instanceof PyFile || node.getElementType() == PyElementTypes.STATEMENT_LIST;
+  }
 }

python/testData/folding/CustomFolding.py

@@ -0,0 +1,25 @@
+<fold text='SupportStuff'>#region SupportStuff
+class SupportClass(object):<fold text='...'>
+    pass</fold>
+
+
+def supportFunction():<fold text='...'>
+    return</fold>
+#endregion</fold>
+
+class ImportantClass(object):<fold text='...'>
+    <fold text='Class body'>#region Class body
+    def f(self):<fold text='...'>
+            return 'hello world'</fold>
+    #endregion</fold>
+    pass</fold>
+
+<fold text='VeryImportant'>#region VeryImportant
+def importantFunction():<fold text='...'>
+    if (x > 0):
+        <fold text='Body'>#region Body
+        print "abc"
+        #endregion</fold>
+    return</fold>
+pass
+#endregion</fold>

python/testSrc/com/jetbrains/python/PyFoldingTest.java

@@ -17,4 +17,8 @@ public class PyFoldingTest extends PyTestCase {
   public void testDocString() {
     doTest();
   }
+
+  public void testCustomFolding() {
+    doTest();
+  }
 }