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 21ba9f03f639..f4f76b3fd515 100644
--- a/platform/lang-api/src/com/intellij/lang/folding/CustomFoldingSurroundDescriptor.java
+++ b/platform/lang-api/src/com/intellij/lang/folding/CustomFoldingSurroundDescriptor.java
@@ -32,6 +32,7 @@ import com.intellij.psi.PsiWhiteSpace;
import com.intellij.psi.codeStyle.CodeStyleManager;
import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
+import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.IncorrectOperationException;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
@@ -73,10 +74,12 @@ public class CustomFoldingSurroundDescriptor implements SurroundDescriptor {
if (startElement != null) {
endElement = findClosestParentBeforeLineBreak(endElement);
if (endElement != null) {
- PsiElement commonParent = startElement.getParent();
+ startElement = adjustStartElementIfEndAbsorbed(startElement, endElement);
+ endElement = adjustEndElementIfStartAbsorbed(startElement, endElement);
+ final PsiElement commonParent = startElement.getParent();
if (endElement.getParent() == commonParent) {
- if (startElement == endElement) return new PsiElement[] {startElement};
- return new PsiElement[] {startElement, endElement};
+ if (startElement == endElement) return new PsiElement[]{startElement};
+ return new PsiElement[]{startElement, endElement};
}
}
}
@@ -84,6 +87,22 @@ public class CustomFoldingSurroundDescriptor implements SurroundDescriptor {
return PsiElement.EMPTY_ARRAY;
}
+ @NotNull
+ private static PsiElement adjustEndElementIfStartAbsorbed(@NotNull PsiElement start, @NotNull PsiElement end) {
+ if (PsiTreeUtil.isAncestor(end, start, false) && start.getTextRange().getEndOffset() == end.getTextRange().getEndOffset()) {
+ return start;
+ }
+ return end;
+ }
+
+ @NotNull
+ private static PsiElement adjustStartElementIfEndAbsorbed(@NotNull PsiElement start, @NotNull PsiElement end) {
+ if (PsiTreeUtil.isAncestor(start, end, false) && start.getTextRange().getStartOffset() == end.getTextRange().getStartOffset()) {
+ return end;
+ }
+ return start;
+ }
+
@Nullable
private static PsiElement findClosestParentAfterLineBreak(PsiElement element) {
PsiElement parent = element;
diff --git a/python/testData/surround/SurroundFirstMethodWithCustomFoldingRegion.py b/python/testData/surround/SurroundFirstMethodWithCustomFoldingRegion.py
new file mode 100644
index 000000000000..0323c905851a
--- /dev/null
+++ b/python/testData/surround/SurroundFirstMethodWithCustomFoldingRegion.py
@@ -0,0 +1,15 @@
+class ThisIsATest():
+ def __init__(self):
+ self.test = 1
+
+ def another_one(self):
+ print "Hello, world!"
+
+ def another_two(self):
+ print "Hello, world!"
+
+ def another_three(self):
+ print "Hello, world!"
+
+ def another_four(self):
+ print "Hello, world!"
\ No newline at end of file
diff --git a/python/testData/surround/SurroundLastMethodWithCustomFoldingRegion.py b/python/testData/surround/SurroundLastMethodWithCustomFoldingRegion.py
new file mode 100644
index 000000000000..2dbc0cb2d23d
--- /dev/null
+++ b/python/testData/surround/SurroundLastMethodWithCustomFoldingRegion.py
@@ -0,0 +1,15 @@
+class ThisIsATest():
+ def __init__(self):
+ self.test = 1
+
+ def another_one(self):
+ print "Hello, world!"
+
+ def another_two(self):
+ print "Hello, world!"
+
+ def another_three(self):
+ print "Hello, world!"
+
+ def another_four(self):
+ print "Hello, world!"
diff --git a/python/testSrc/com/jetbrains/python/PySurroundWithTest.java b/python/testSrc/com/jetbrains/python/PySurroundWithTest.java
index 5f49581f1e71..fc7136f3bf52 100644
--- a/python/testSrc/com/jetbrains/python/PySurroundWithTest.java
+++ b/python/testSrc/com/jetbrains/python/PySurroundWithTest.java
@@ -16,9 +16,14 @@
package com.jetbrains.python;
import com.intellij.codeInsight.generation.surroundWith.SurroundWithHandler;
+import com.intellij.lang.folding.CustomFoldingSurroundDescriptor;
import com.intellij.lang.surroundWith.Surrounder;
import com.intellij.openapi.command.WriteCommandAction;
+import com.intellij.openapi.editor.SelectionModel;
+import com.intellij.psi.PsiElement;
import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyFunction;
import com.jetbrains.python.refactoring.surround.surrounders.statements.PyWithIfSurrounder;
import com.jetbrains.python.refactoring.surround.surrounders.statements.PyWithTryExceptSurrounder;
import com.jetbrains.python.refactoring.surround.surrounders.statements.PyWithWhileSurrounder;
@@ -39,6 +44,29 @@ public class PySurroundWithTest extends PyTestCase {
doTest(new PyWithTryExceptSurrounder());
}
+ // PY-11357
+ public void testSurroundFirstMethodWithCustomFoldingRegion() {
+ checkCustomFoldingRegionRange(PyFunction.class);
+ }
+
+ // PY-11357
+ public void testSurroundLastMethodWithCustomFoldingRegion() {
+ checkCustomFoldingRegionRange(PyFunction.class);
+ }
+
+ private PsiElement[] checkCustomFoldingRegionRange(Class extends PyElement>... elementTypes) {
+ myFixture.configureByFile("/surround/" + getTestName(false) + ".py");
+ final SelectionModel selection = myFixture.getEditor().getSelectionModel();
+ final PsiElement[] range = CustomFoldingSurroundDescriptor.INSTANCE.getElementsToSurround(myFixture.getFile(),
+ selection.getSelectionStart(),
+ selection.getSelectionEnd());
+ assertEquals(elementTypes.length, range.length);
+ for (int i = 0; i < elementTypes.length; i++) {
+ assertInstanceOf(range[i], elementTypes[i]);
+ }
+ return range;
+ }
+
private void doTest(final Surrounder surrounder) throws Exception {
String baseName = "/surround/" + getTestName(false);
myFixture.configureByFile(baseName + ".py");