PY-5558 Provide folding for if/while and other code blocks

python/src/com/jetbrains/python/PyElementTypes.java

@@ -146,7 +146,7 @@ public interface PyElementTypes {
 
   PyElementType ELSE_PART = new PyElementType("ELSE_PART", PyElsePartImpl.class);
 
-  TokenSet PARTS = TokenSet.create(IF_PART_IF, IF_PART_ELIF, FOR_PART, WHILE_PART, TRY_PART, FINALLY_PART, ELSE_PART); 
+  TokenSet PARTS = TokenSet.create(IF_PART_IF, IF_PART_ELIF, FOR_PART, WHILE_PART, TRY_PART, FINALLY_PART, ELSE_PART, EXCEPT_PART);
   TokenSet ELIFS = TokenSet.create(IF_PART_ELIF);
   TokenSet STAR_PARAMETERS = TokenSet.create(NAMED_PARAMETER, STAR_ARGUMENT_EXPRESSION, STAR_EXPRESSION, DOUBLE_STAR_EXPRESSION);
   TokenSet CLASS_OR_FUNCTION = TokenSet.create(CLASS_DECLARATION, FUNCTION_DECLARATION);

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

@@ -26,10 +26,7 @@ import com.intellij.openapi.util.text.LineTokenizer;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.tree.IElementType;
-import com.jetbrains.python.psi.PyFile;
-import com.jetbrains.python.psi.PyFileElementType;
-import com.jetbrains.python.psi.PyImportStatementBase;
-import com.jetbrains.python.psi.PyStringLiteralExpression;
+import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFileImpl;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -75,7 +72,9 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw
 
   private static void foldStatementList(ASTNode node, List<FoldingDescriptor> descriptors) {
     IElementType elType = node.getTreeParent().getElementType();
-    if (elType == PyElementTypes.FUNCTION_DECLARATION || elType == PyElementTypes.CLASS_DECLARATION) {
+    if (elType == PyElementTypes.FUNCTION_DECLARATION
+        || elType == PyElementTypes.CLASS_DECLARATION
+        || ifFoldBlocks(node, elType)) {
       ASTNode colon = node.getTreeParent().findChildByType(PyTokenTypes.COLON);
       if (colon != null && colon.getStartOffset() + 1 < node.getTextRange().getEndOffset() - 1) {
         final CharSequence chars = node.getChars();
@@ -95,6 +94,18 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw
     }
   }
 
+  private static boolean ifFoldBlocks(ASTNode statementList, IElementType parentType) {
+    if (!PyElementTypes.PARTS.contains(parentType)) {
+      return false;
+    }
+    PsiElement element = statementList.getPsi();
+    if (element instanceof PyStatementList) {
+      PyStatementList statements = (PyStatementList)element;
+      return statements.getStatements().length > 1;
+    }
+    return false;
+  }
+
   private static void foldDocString(ASTNode node, List<FoldingDescriptor> descriptors) {
     if (getDocStringOwnerType(node) != null && StringUtil.countChars(node.getText(), '\n') > 1) {
       descriptors.add(new FoldingDescriptor(node, node.getTextRange()));

python/testData/folding/blocksFolding.py

@@ -0,0 +1,51 @@
+if True:
+    pass
+else:
+    pass
+
+if True:<fold text='...'>
+    pass
+    pass</fold>
+elif True:<fold text='...'>
+    pass
+    pass</fold>
+else:<fold text='...'>
+    pass
+    pass</fold>
+
+
+x  = []
+for i in x:
+    pass
+
+for i in x:<fold text='...'>
+    pass
+    pass</fold>
+
+while True:<fold text='...'>
+    pass
+    pass</fold>
+
+f = open('1.txt')
+ints = []
+try:
+    for line in f:<fold text='...'>
+        ints.append(int(line))
+        ints.append(int(line))</fold>
+
+except ValueError:<fold text='...'>
+    print('')
+    print('')
+    print('')
+    print('')
+    print('')</fold>
+except Exception:
+    print('')
+else:<fold text='...'>
+    print('')
+    print('')
+    print('')
+    print('')</fold>
+finally:<fold text='...'>
+    f.close()
+    print('')</fold>

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

@@ -40,4 +40,8 @@ public class PyFoldingTest extends PyTestCase {
   public void testImportBlock() {
     doTest();
   }
+
+  public void testBlocksFolding() {
+    doTest();
+  }
 }