From f8aecd088e170fb3eb6a46c9f2708271b2748448 Mon Sep 17 00:00:00 2001 From: Mikhail Golubev Date: Fri, 19 Nov 2021 17:55:09 +0200 Subject: [PATCH] EA-401551 Fix an exception when joining IFs with assignment expressions GitOrigin-RevId: ce9610afb4e280b2930f126861bdaf13bd69021e --- .../intentions/PyJoinIfIntention.java | 17 +++++++++++++---- .../refactoring/PyReplaceExpressionUtil.java | 3 ++- ...oinIfAssignmentExpressionInInnerCondition.py | 3 +++ ...ssignmentExpressionInInnerCondition_after.py | 2 ++ ...inIfAssignmentExpressionsInBothConditions.py | 3 +++ ...signmentExpressionsInBothConditions_after.py | 2 ++ .../joinIfOrExpressionInOuterCondition.py | 3 +++ .../joinIfOrExpressionInOuterCondition_after.py | 2 ++ .../python/intentions/PyIntentionTest.java | 13 +++++++++++++ 9 files changed, 43 insertions(+), 5 deletions(-) create mode 100644 python/testData/intentions/joinIfAssignmentExpressionInInnerCondition.py create mode 100644 python/testData/intentions/joinIfAssignmentExpressionInInnerCondition_after.py create mode 100644 python/testData/intentions/joinIfAssignmentExpressionsInBothConditions.py create mode 100644 python/testData/intentions/joinIfAssignmentExpressionsInBothConditions_after.py create mode 100644 python/testData/intentions/joinIfOrExpressionInOuterCondition.py create mode 100644 python/testData/intentions/joinIfOrExpressionInOuterCondition_after.py diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/intentions/PyJoinIfIntention.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/intentions/PyJoinIfIntention.java index 6c76d17db8bd..312528f41344 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/intentions/PyJoinIfIntention.java +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/intentions/PyJoinIfIntention.java @@ -10,8 +10,8 @@ import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.ArrayUtil; import com.intellij.util.IncorrectOperationException; import com.jetbrains.python.PyPsiBundle; -import com.jetbrains.python.PyTokenTypes; import com.jetbrains.python.psi.*; +import com.jetbrains.python.refactoring.PyReplaceExpressionUtil; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -81,15 +81,24 @@ public class PyJoinIfIntention extends PyBaseIntentionAction { if (outerCondition == null || innerCondition == null) return; PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); - StringBuilder replacementText = new StringBuilder(outerCondition.getText() + " and "); - if (innerCondition instanceof PyBinaryExpression && ((PyBinaryExpression)innerCondition).getOperator() == PyTokenTypes.OR_KEYWORD) { + LanguageLevel pyVersion = LanguageLevel.forElement(file); + PyBinaryExpression fakeAndExpression = (PyBinaryExpression)elementGenerator.createExpressionFromText(pyVersion, "foo and bar"); + StringBuilder replacementText = new StringBuilder(); + if (PyReplaceExpressionUtil.isNeedParenthesis(fakeAndExpression.getLeftExpression(), outerCondition)) { + replacementText.append("(").append(outerCondition.getText()).append(")"); + } + else { + replacementText.append(outerCondition.getText()); + } + replacementText.append(" and "); + if (PyReplaceExpressionUtil.isNeedParenthesis(fakeAndExpression.getLeftExpression(), innerCondition)) { replacementText.append("(").append(innerCondition.getText()).append(")"); } else { replacementText.append(innerCondition.getText()); } - PyExpression newCondition = elementGenerator.createExpressionFromText(LanguageLevel.forElement(file), replacementText.toString()); + PyExpression newCondition = elementGenerator.createExpressionFromText(pyVersion, replacementText.toString()); outerCondition.replace(newCondition); PyStatementList innerStatementList = innerIfStatement.getIfPart().getStatementList(); diff --git a/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyReplaceExpressionUtil.java b/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyReplaceExpressionUtil.java index f13de9c479e4..7a3446e29a52 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyReplaceExpressionUtil.java +++ b/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyReplaceExpressionUtil.java @@ -433,7 +433,7 @@ public final class PyReplaceExpressionUtil implements PyElementTypes { opType == DIV || opType == FLOORDIV || opType == PERC || opType == EXP || opType == MINUS; } - private static int getExpressionPriority(PyElement expr) { + private static int getExpressionPriority(@NotNull PyElement expr) { int priority = 0; if (expr instanceof PyReferenceExpression || expr instanceof PySubscriptionExpression || @@ -459,6 +459,7 @@ public final class PyReplaceExpressionUtil implements PyElementTypes { } else if (expr instanceof PyConditionalExpression) priority = 14; else if (expr instanceof PyLambdaExpression) priority = 15; + else if (expr instanceof PyAssignmentExpression) priority = 16; return -priority; } diff --git a/python/testData/intentions/joinIfAssignmentExpressionInInnerCondition.py b/python/testData/intentions/joinIfAssignmentExpressionInInnerCondition.py new file mode 100644 index 000000000000..b8349ba8148b --- /dev/null +++ b/python/testData/intentions/joinIfAssignmentExpressionInInnerCondition.py @@ -0,0 +1,3 @@ +if x: + if y := foo(): + pass diff --git a/python/testData/intentions/joinIfAssignmentExpressionInInnerCondition_after.py b/python/testData/intentions/joinIfAssignmentExpressionInInnerCondition_after.py new file mode 100644 index 000000000000..ca63164ec861 --- /dev/null +++ b/python/testData/intentions/joinIfAssignmentExpressionInInnerCondition_after.py @@ -0,0 +1,2 @@ +if x and (y := foo()): + pass diff --git a/python/testData/intentions/joinIfAssignmentExpressionsInBothConditions.py b/python/testData/intentions/joinIfAssignmentExpressionsInBothConditions.py new file mode 100644 index 000000000000..1c16ba6440c2 --- /dev/null +++ b/python/testData/intentions/joinIfAssignmentExpressionsInBothConditions.py @@ -0,0 +1,3 @@ +if x := foo(): + if y := bar(): + pass diff --git a/python/testData/intentions/joinIfAssignmentExpressionsInBothConditions_after.py b/python/testData/intentions/joinIfAssignmentExpressionsInBothConditions_after.py new file mode 100644 index 000000000000..393ad7409990 --- /dev/null +++ b/python/testData/intentions/joinIfAssignmentExpressionsInBothConditions_after.py @@ -0,0 +1,2 @@ +if (x := foo()) and (y := bar()): + pass diff --git a/python/testData/intentions/joinIfOrExpressionInOuterCondition.py b/python/testData/intentions/joinIfOrExpressionInOuterCondition.py new file mode 100644 index 000000000000..55796e92b567 --- /dev/null +++ b/python/testData/intentions/joinIfOrExpressionInOuterCondition.py @@ -0,0 +1,3 @@ +if x or y: + if z: + pass diff --git a/python/testData/intentions/joinIfOrExpressionInOuterCondition_after.py b/python/testData/intentions/joinIfOrExpressionInOuterCondition_after.py new file mode 100644 index 000000000000..f8667e21a3cd --- /dev/null +++ b/python/testData/intentions/joinIfOrExpressionInOuterCondition_after.py @@ -0,0 +1,2 @@ +if (x or y) and z: + pass diff --git a/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java index 71fad8cd04d1..0ab09a755f5f 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java @@ -153,6 +153,19 @@ public class PyIntentionTest extends PyTestCase { doNegativeTest(PyPsiBundle.message("INTN.join.if")); } + public void testJoinIfOrExpressionInOuterCondition() { + doTest(PyPsiBundle.message("INTN.join.if")); + } + + // EA-401551 + public void testJoinIfAssignmentExpressionInInnerCondition() { + doTest(PyPsiBundle.message("INTN.join.if")); + } + + public void testJoinIfAssignmentExpressionsInBothConditions() { + doTest(PyPsiBundle.message("INTN.join.if")); + } + public void testDictConstructorToLiteralForm() { doTest(PyPsiBundle.message("INTN.convert.dict.constructor.to.dict.literal")); }