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... 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");