PY-14261 Surrounding with custom folding preserves indentation in Python

platform/lang-api/src/com/intellij/lang/folding/CustomFoldingSurroundDescriptor.java

@@ -21,6 +21,7 @@ import com.intellij.lang.Language;
 import com.intellij.lang.LanguageCommenters;
 import com.intellij.lang.surroundWith.SurroundDescriptor;
 import com.intellij.lang.surroundWith.Surrounder;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
@@ -198,6 +199,9 @@ public class CustomFoldingSurroundDescriptor implements SurroundDescriptor {
       if (linePrefix == null) return null;
       int prefixLength = linePrefix.length();
       int startOffset = firstElement.getTextRange().getStartOffset();
+      final Document document = editor.getDocument();
+      final int startLineNumber = document.getLineNumber(startOffset);
+      final String startIndent = document.getText(new TextRange(document.getLineStartOffset(startLineNumber), startOffset));
       int endOffset = lastElement.getTextRange().getEndOffset();
       int delta = 0;
       TextRange rangeToSelect = new TextRange(startOffset, startOffset);
@@ -207,11 +211,11 @@ public class CustomFoldingSurroundDescriptor implements SurroundDescriptor {
         startText = startText.replace("?", DEFAULT_DESC_TEXT);
         rangeToSelect = new TextRange(startOffset + descPos, startOffset + descPos + DEFAULT_DESC_TEXT.length());
       }
-      String startString = linePrefix + startText + "\n";
+      String startString = linePrefix + startText + "\n" + startIndent;
       String endString = "\n" + linePrefix + myProvider.getEndString();
-      editor.getDocument().insertString(endOffset, endString);
+      document.insertString(endOffset, endString);
       delta += endString.length();
-      editor.getDocument().insertString(startOffset, startString);
+      document.insertString(startOffset, startString);
       delta += startString.length();
       rangeToSelect = rangeToSelect.shiftRight(prefixLength);
       PsiDocumentManager documentManager = PsiDocumentManager.getInstance(project);

python/testData/surround/SurroundWithCustomFoldingRegion.py

@@ -0,0 +1,5 @@
+def func(x):
+    assert x is not None
+    <selection>if True:
+        pass</selection>
+    return x

python/testData/surround/SurroundWithCustomFoldingRegion_after.py

@@ -0,0 +1,7 @@
+def func(x):
+    assert x is not None
+    # <editor-fold desc="Description">
+    if True:
+        pass
+    # </editor-fold>
+    return x

python/testSrc/com/jetbrains/python/PySurroundWithTest.java

@@ -54,6 +54,12 @@ public class PySurroundWithTest extends PyTestCase {
     checkCustomFoldingRegionRange(PyFunction.class);
   }
 
+  // PY-14261
+  public void testSurroundWithCustomFoldingRegion() throws Exception {
+    doTest(CustomFoldingSurroundDescriptor.SURROUNDERS[0]);
+  }
+
+
   private PsiElement[] checkCustomFoldingRegionRange(Class<? extends PyElement>... elementTypes) {
     myFixture.configureByFile("/surround/" + getTestName(false) + ".py");
     final SelectionModel selection = myFixture.getEditor().getSelectionModel();