From 9a15bf3b54dfdeb19d2ed404e5883c00de67577e Mon Sep 17 00:00:00 2001
From: Liana Bakradze <liana.bakradze@jetbrains.com>
Date: Wed, 11 Nov 2015 15:59:49 +0300
Subject: [PATCH] PY-5558 Provide folding for if/while and other code blocks

---
 .../com/jetbrains/python/PyElementTypes.java  |  2 +-
 .../python/PythonFoldingBuilder.java          | 21 ++++++--
 python/testData/folding/blocksFolding.py      | 51 +++++++++++++++++++
 .../com/jetbrains/python/PyFoldingTest.java   |  4 ++
 4 files changed, 72 insertions(+), 6 deletions(-)
 create mode 100644 python/testData/folding/blocksFolding.py

diff --git a/python/src/com/jetbrains/python/PyElementTypes.java b/python/src/com/jetbrains/python/PyElementTypes.java
index 5e76006fd940..4d7c74fc85ba 100644
--- a/python/src/com/jetbrains/python/PyElementTypes.java
+++ b/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);
diff --git a/python/src/com/jetbrains/python/PythonFoldingBuilder.java b/python/src/com/jetbrains/python/PythonFoldingBuilder.java
index 07cafac2dd93..4127806be73b 100644
--- a/python/src/com/jetbrains/python/PythonFoldingBuilder.java
+++ b/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()));
diff --git a/python/testData/folding/blocksFolding.py b/python/testData/folding/blocksFolding.py
new file mode 100644
index 000000000000..d55f3865516d
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyFoldingTest.java b/python/testSrc/com/jetbrains/python/PyFoldingTest.java
index 5b66747fb9db..201376efcd9f 100644
--- a/python/testSrc/com/jetbrains/python/PyFoldingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyFoldingTest.java
@@ -40,4 +40,8 @@ public class PyFoldingTest extends PyTestCase {
   public void testImportBlock() {
     doTest();
   }
+
+  public void testBlocksFolding() {
+    doTest();
+  }
 }