diff --git a/platform/lang-api/src/com/intellij/lang/folding/CustomFoldingSurroundDescriptor.java b/platform/lang-api/src/com/intellij/lang/folding/CustomFoldingSurroundDescriptor.java
index f4f76b3fd515..17da710c06c5 100644
--- a/platform/lang-api/src/com/intellij/lang/folding/CustomFoldingSurroundDescriptor.java
+++ b/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);
diff --git a/python/testData/surround/SurroundWithCustomFoldingRegion.py b/python/testData/surround/SurroundWithCustomFoldingRegion.py
new file mode 100644
index 000000000000..3fb0cdb75d5b
--- /dev/null
+++ b/python/testData/surround/SurroundWithCustomFoldingRegion.py
@@ -0,0 +1,5 @@
+def func(x):
+ assert x is not None
+ if True:
+ pass
+ return x
\ No newline at end of file
diff --git a/python/testData/surround/SurroundWithCustomFoldingRegion_after.py b/python/testData/surround/SurroundWithCustomFoldingRegion_after.py
new file mode 100644
index 000000000000..d1c19a492e54
--- /dev/null
+++ b/python/testData/surround/SurroundWithCustomFoldingRegion_after.py
@@ -0,0 +1,7 @@
+def func(x):
+ assert x is not None
+ #
+ if True:
+ pass
+ #
+ return x
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PySurroundWithTest.java b/python/testSrc/com/jetbrains/python/PySurroundWithTest.java
index fc7136f3bf52..9abfa6313669 100644
--- a/python/testSrc/com/jetbrains/python/PySurroundWithTest.java
+++ b/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();