From 2bf033e420e671f9ec96440a788b490353068a2b Mon Sep 17 00:00:00 2001
From: somethingnew179 <kondakov.ilya1@gmail.com>
Date: Fri, 19 Jul 2024 22:46:34 +0300
Subject: [PATCH] PY-49174 make Match and Case expressions folding

GitOrigin-RevId: c9b1b37a762ec0879eb2b7653ae14b34ae97dfc0
---
 .../python/PythonFoldingBuilder.java          | 56 +++++++++++++------
 python/testData/folding/matchFolding.py       | 19 +++++++
 .../com/jetbrains/python/PyFoldingTest.java   |  5 ++
 3 files changed, 62 insertions(+), 18 deletions(-)
 create mode 100644 python/testData/folding/matchFolding.py

diff --git a/python/src/com/jetbrains/python/PythonFoldingBuilder.java b/python/src/com/jetbrains/python/PythonFoldingBuilder.java
index fb56c09917fc..63a596427dc8 100644
--- a/python/src/com/jetbrains/python/PythonFoldingBuilder.java
+++ b/python/src/com/jetbrains/python/PythonFoldingBuilder.java
@@ -56,6 +56,9 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw
                                                                          lastImport.getTextRange().getEndOffset())));
       }
     }
+    else if (elementType == PyElementTypes.MATCH_STATEMENT) {
+      foldMatchStatement(node, descriptors);
+    }
     else if (elementType == PyElementTypes.STATEMENT_LIST) {
       foldStatementList(node, descriptors);
     }
@@ -129,6 +132,36 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw
     }
   }
 
+  private static void foldMatchStatement(@NotNull ASTNode node, @NotNull List<FoldingDescriptor> descriptors) {
+    TextRange nodeRange = node.getTextRange();
+    if (nodeRange.isEmpty()) {
+      return;
+    }
+
+    IElementType elType = node.getElementType();
+    if (elType == PyElementTypes.MATCH_STATEMENT) {
+      ASTNode colon = node.findChildByType(PyTokenTypes.COLON);
+      foldSegment(node, descriptors, nodeRange, colon);
+    }
+  }
+
+  private static void foldSegment(@NotNull ASTNode node, @NotNull List<FoldingDescriptor> descriptors, @NotNull TextRange nodeRange, @Nullable ASTNode colon) {
+    int nodeEnd = nodeRange.getEndOffset();
+    if (colon != null && nodeEnd - (colon.getStartOffset() + 1) > 1) {
+      CharSequence chars = node.getChars();
+      int nodeStart = nodeRange.getStartOffset();
+      int foldStart = colon.getStartOffset() + 1;
+      int foldEnd = nodeEnd;
+      while (foldEnd > Math.max(nodeStart, foldStart + 1) && Character.isWhitespace(chars.charAt(foldEnd - nodeStart - 1))) {
+        foldEnd--;
+      }
+      descriptors.add(new FoldingDescriptor(node, new TextRange(foldStart, foldEnd)));
+    }
+    else if (nodeRange.getLength() > 1) { // only for ranges at least 1 char wide
+      descriptors.add(new FoldingDescriptor(node, nodeRange));
+    }
+  }
+
   private static void foldStatementList(ASTNode node, List<FoldingDescriptor> descriptors) {
     final TextRange nodeRange = node.getTextRange();
     if (nodeRange.isEmpty()) {
@@ -136,27 +169,14 @@ public class PythonFoldingBuilder extends CustomFoldingBuilder implements DumbAw
     }
 
     final IElementType elType = node.getTreeParent().getElementType();
-    if (elType == PyElementTypes.FUNCTION_DECLARATION || elType == PyElementTypes.CLASS_DECLARATION || ifFoldBlocks(node, elType)) {
-      final ASTNode colon = node.getTreeParent().findChildByType(PyTokenTypes.COLON);
-      final int nodeEnd = nodeRange.getEndOffset();
-      if (colon != null && nodeEnd - (colon.getStartOffset() + 1) > 1) {
-        final CharSequence chars = node.getChars();
-        final int nodeStart = nodeRange.getStartOffset();
-        final int foldStart = colon.getStartOffset() + 1;
-        int foldEnd = nodeEnd;
-        while (foldEnd > Math.max(nodeStart, foldStart + 1) && Character.isWhitespace(chars.charAt(foldEnd - nodeStart - 1))) {
-          foldEnd--;
-        }
-        descriptors.add(new FoldingDescriptor(node, new TextRange(foldStart, foldEnd)));
-      }
-      else if (nodeRange.getLength() > 1) { // only for ranges at least 1 char wide
-        descriptors.add(new FoldingDescriptor(node, nodeRange));
-      }
+    if (elType == PyElementTypes.FUNCTION_DECLARATION || elType == PyElementTypes.CLASS_DECLARATION || checkFoldBlocks(node, elType)) {
+      ASTNode colon = node.getTreeParent().findChildByType(PyTokenTypes.COLON);
+      foldSegment(node, descriptors, nodeRange, colon);
     }
   }
 
-  private static boolean ifFoldBlocks(ASTNode statementList, IElementType parentType) {
-    if (!PyElementTypes.PARTS.contains(parentType) && parentType != PyElementTypes.WITH_STATEMENT) {
+  private static boolean checkFoldBlocks(@NotNull ASTNode statementList, @NotNull IElementType parentType) {
+    if (!PyElementTypes.PARTS.contains(parentType) && parentType != PyElementTypes.WITH_STATEMENT && parentType != PyElementTypes.CASE_CLAUSE) {
       return false;
     }
     PsiElement element = statementList.getPsi();
diff --git a/python/testData/folding/matchFolding.py b/python/testData/folding/matchFolding.py
new file mode 100644
index 000000000000..90a0b3572d6f
--- /dev/null
+++ b/python/testData/folding/matchFolding.py
@@ -0,0 +1,19 @@
+    def f(data):<fold text='...'>
+        match data:<fold text='...'>
+            case 1:<fold text='...'>
+                pass
+                pass</fold>
+            case 2:
+                pass
+            case 3:<fold text='...'>
+                pass
+                pass</fold></fold></fold>
+
+
+    p = 0
+    match p:<fold text='...'>
+        case 0:<fold text='...'>
+            pass
+            pass</fold>
+        case 1:
+            pass</fold>
diff --git a/python/testSrc/com/jetbrains/python/PyFoldingTest.java b/python/testSrc/com/jetbrains/python/PyFoldingTest.java
index 8e9b68e0c3ac..729387ddec36 100644
--- a/python/testSrc/com/jetbrains/python/PyFoldingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyFoldingTest.java
@@ -101,4 +101,9 @@ public class PyFoldingTest extends PyTestCase {
   public void testStringPrefixFolding() {
     doTest();
   }
+
+  // PY-49174
+  public void testMatchFolding() {
+    doTest();
+  }
 }