PY-33843 Fix problem with extracting constant from comprehension

(cherry picked from commit 72bb8e523471c9922ae3d8f95f7fa9ce1fe0009d)

IJ-MR-18836

GitOrigin-RevId: dadb7cf489a0501b834cff7e10979fab2c8ae8cd

python/src/com/jetbrains/python/refactoring/introduce/IntroduceHandler.java

@@ -41,6 +41,7 @@ import com.intellij.refactoring.listeners.RefactoringEventData;
 import com.intellij.refactoring.listeners.RefactoringEventListener;
 import com.intellij.refactoring.util.CommonRefactoringUtil;
 import com.intellij.util.ObjectUtils;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
@@ -48,6 +49,7 @@ import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
+import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import com.jetbrains.python.psi.types.PyCallableParameter;
 import com.jetbrains.python.psi.types.PyNoneType;
 import com.jetbrains.python.psi.types.PyType;
@@ -295,8 +297,7 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
     }
 
     element1 = PyRefactoringUtil.getSelectedExpression(project, file, element1, element2);
-    final PyComprehensionElement comprehension = PsiTreeUtil.getParentOfType(element1, PyComprehensionElement.class, true);
-    if (element1 == null || comprehension != null) {
+    if (element1 == null || referencesComprehensionIteratorValue(element1)) {
       showCannotPerformError(project, editor);
       return;
     }
@@ -333,6 +334,45 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
     performActionOnElement(operation);
   }
 
+  private static boolean referencesComprehensionIteratorValue(@NotNull PsiElement element) {
+    PyComprehensionElement comprehension = PsiTreeUtil.getParentOfType(element, PyComprehensionElement.class, true);
+    if (comprehension != null) {
+      List<PyExpression> iteratorVariables = ContainerUtil.map(comprehension.getForComponents(), it -> it.getIteratorVariable());
+      List<PsiElement> referencedInSelection = collectReferencedDefinitionsInSameFile(element, element.getContainingFile());
+      if (iteratorVariables.contains(element) || ContainerUtil.intersects(referencedInSelection, iteratorVariables)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  protected static @NotNull List<PsiElement> collectReferencedDefinitionsInSameFile(@NotNull PsiElement element, @NotNull PsiFile file) {
+    PsiElement selectionElement = getOriginalSelectionCoveringElement(element);
+    TextRange textRange = getTextRangeForOperationElement(element);
+
+    return StreamEx.of(PsiTreeUtil.collectElementsOfType(selectionElement, PyReferenceExpression.class))
+      .filter(it -> textRange.contains(it.getTextRange()))
+      .filter(ref -> !ref.isQualified())
+      .flatMap(expr -> PyResolveUtil.resolveLocally(expr).stream())
+      .filter(it -> it != null && it.getContainingFile() == file)
+      .toList();
+  }
+
+  protected static @NotNull TextRange getTextRangeForOperationElement(@NotNull PsiElement operationElement) {
+    var userData = operationElement.getUserData(PyReplaceExpressionUtil.SELECTION_BREAKS_AST_NODE);
+    if (userData == null || userData.first == null || userData.second == null) {
+      return operationElement.getTextRange();
+    }
+    else {
+      return userData.second.shiftRight(userData.first.getTextOffset());
+    }
+  }
+
+  protected static @NotNull PsiElement getOriginalSelectionCoveringElement(@NotNull PsiElement operationElement) {
+    var userData = operationElement.getUserData(PyReplaceExpressionUtil.SELECTION_BREAKS_AST_NODE);
+    return userData == null ? operationElement : userData.first;
+  }
+
   private static boolean breaksStringFormatting(@NotNull String s, @NotNull TextRange range) {
     return breaksRanges(substitutionsToRanges(filterSubstitutions(parsePercentFormat(s))), range);
   }
@@ -421,8 +461,7 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
     if (call != null && call.getCallee() == element) {
       return false;
     }
-    final PyComprehensionElement comprehension = PsiTreeUtil.getParentOfType(element, PyComprehensionElement.class, true);
-    if (comprehension != null) {
+    if (referencesComprehensionIteratorValue(element)) {
       return false;
     }
     return true;

python/src/com/jetbrains/python/refactoring/introduce/constant/PyIntroduceConstantHandler.java

@@ -67,7 +67,7 @@ public class PyIntroduceConstantHandler extends IntroduceHandler {
     assert containingFile instanceof PyFile;
     PsiElement initialPosition = AddImportHelper.getFileInsertPosition((PyFile)containingFile);
 
-    List<PsiElement> sameFileRefs = collectReferencedDefinitionsInSameFile(operation);
+    List<PsiElement> sameFileRefs = collectReferencedDefinitionsInSameFile(operation.getElement(), operation.getFile());
     PsiElement maxPosition = getLowermostTopLevelStatement(sameFileRefs);
 
     if (maxPosition == null) {
@@ -102,7 +102,7 @@ public class PyIntroduceConstantHandler extends IntroduceHandler {
     Editor editor = operation.getEditor();
     if (editor == null) return false;
 
-    List<PsiElement> sameFileRefs = collectReferencedDefinitionsInSameFile(operation);
+    List<PsiElement> sameFileRefs = collectReferencedDefinitionsInSameFile(operation.getElement(), operation.getFile());
     if (!ContainerUtil.all(sameFileRefs, it -> PyUtil.isTopLevel(it))) return false;
     PsiElement maxPosition = getLowermostTopLevelStatement(sameFileRefs);
     if (maxPosition == null) return true;
@@ -110,33 +110,6 @@ public class PyIntroduceConstantHandler extends IntroduceHandler {
            !PsiTreeUtil.isAncestor(maxPosition, selectionElement, false);
   }
 
-  private static @NotNull List<PsiElement> collectReferencedDefinitionsInSameFile(@NotNull IntroduceOperation operation) {
-    PsiElement selectionElement = getOriginalSelectionCoveringElement(operation.getElement());
-    TextRange textRange = getTextRangeForOperationElement(operation.getElement());
-
-    return StreamEx.of(PsiTreeUtil.collectElementsOfType(selectionElement, PyReferenceExpression.class))
-      .filter(it -> textRange.contains(it.getTextRange()))
-      .filter(ref -> !ref.isQualified())
-      .flatMap(expr -> PyResolveUtil.resolveLocally(expr).stream())
-      .filter(it -> it != null && it.getContainingFile() == operation.getFile())
-      .toList();
-  }
-
-  private static @NotNull TextRange getTextRangeForOperationElement(@NotNull PsiElement operationElement) {
-    var userData = operationElement.getUserData(PyReplaceExpressionUtil.SELECTION_BREAKS_AST_NODE);
-    if (userData == null || userData.first == null || userData.second == null) {
-      return operationElement.getTextRange();
-    }
-    else {
-      return userData.second.shiftRight(userData.first.getTextOffset());
-    }
-  }
-
-  private static @NotNull PsiElement getOriginalSelectionCoveringElement(@NotNull PsiElement operationElement) {
-    var userData = operationElement.getUserData(PyReplaceExpressionUtil.SELECTION_BREAKS_AST_NODE);
-    return userData == null ? operationElement : userData.first;
-  }
-
   private static @Nullable PsiElement getLowermostTopLevelStatement(@NotNull List<PsiElement> elements) {
     return StreamEx.of(elements)
       .map(it -> PyPsiUtils.getParentRightBefore(it, it.getContainingFile()))

python/testData/refactoring/introduceConstant/stringLiteralArgumentInComprehension.after.py

@@ -0,0 +1,5 @@
+a = 'constant'
+
+
+def func(xs):
+    return [x for x in xs if x.startswith(a)]

python/testData/refactoring/introduceConstant/stringLiteralArgumentInComprehension.py

@@ -0,0 +1,2 @@
+def func(xs):
+    return [x for x in xs if x.startswith(<selection>'constant'</selection>)]

python/testData/refactoring/introduceVariable/expressionInComprehensionIteratedList.after.py

@@ -0,0 +1,3 @@
+def foo():
+    a = range(0, 42)<caret>
+    return [x + 1 for x in a]

python/testData/refactoring/introduceVariable/expressionInComprehensionIteratedList.py

@@ -0,0 +1,2 @@
+def foo():
+    return [x + 1 for x in <selection>range(0, 42)</selection>]

python/testData/refactoring/introduceVariable/iteratorDependentExpressionInComprehensionIfComponent.py

@@ -0,0 +1,2 @@
+def foo():
+    return [x + 1 for x in range(0, 42) if <selection>x > 3</selection>]

python/testData/refactoring/introduceVariable/iteratorDependentExpressionInComprehensionResult.py

@@ -0,0 +1,2 @@
+def foo():
+    return [<selection>x + 1</selection> for x in range(0, 42)]

python/testData/refactoring/introduceVariable/iteratorFreeExpressionInComprehensionIfComponent.after.py

@@ -0,0 +1,5 @@
+y = 42
+
+def foo():
+    a = y + 3<caret>
+    return [x + 1 for x in range(0, 42) if x > a]

python/testData/refactoring/introduceVariable/iteratorFreeExpressionInComprehensionIfComponent.py

@@ -0,0 +1,4 @@
+y = 42
+
+def foo():
+    return [x + 1 for x in range(0, 42) if x > <selection>y + 3</selection>]

python/testData/refactoring/introduceVariable/iteratorFreeExpressionInComprehensionResult.after.py

@@ -0,0 +1,5 @@
+y = 42
+
+def foo():
+    a = y + 1<caret>
+    return [x + a for x in range(0, 42)]

python/testData/refactoring/introduceVariable/iteratorFreeExpressionInComprehensionResult.py

@@ -0,0 +1,4 @@
+y = 42
+
+def foo():
+    return [x + <selection>y + 1</selection> for x in range(0, 42)]

python/testData/refactoring/introduceVariable/iteratorVariableDependentExpressionInComprehensionIteratedList.py

@@ -0,0 +1,2 @@
+def foo():
+    return [x + y for x in range(42) for y in <selection>range(x)</selection>]

python/testData/refactoring/introduceVariable/iteratorVariableInComprehension.py

@@ -0,0 +1,2 @@
+def foo():
+    return [x + 1 for <selection>x</selection> in range(0, 42)]

python/testSrc/com/jetbrains/python/refactoring/PyIntroduceConstantTest.java

@@ -119,6 +119,11 @@ public class PyIntroduceConstantTest extends PyIntroduceTestCase {
     doTestThrowsRefactoringErrorHintException();
   }
 
+  // PY-33843
+  public void testStringLiteralArgumentInComprehension() {
+    doTest();
+  }
+
   private void doTestThrowsRefactoringErrorHintException() {
     assertThrows(CommonRefactoringUtil.RefactoringErrorHintException.class, () -> doTest());
   }

python/testSrc/com/jetbrains/python/refactoring/PyIntroduceVariableTest.java

@@ -262,6 +262,41 @@ public class PyIntroduceVariableTest extends PyIntroduceTestCase {
     runWithLanguageLevel(LanguageLevel.PYTHON27, this::doTest);
   }
 
+  // PY-33843
+  public void testIteratorVariableInComprehension() {
+    doTestCannotPerform();
+  }
+
+  // PY-33843
+  public void testIteratorDependentExpressionInComprehensionIfComponent() {
+    doTestCannotPerform();
+  }
+
+  // PY-33843
+  public void testIteratorDependentExpressionInComprehensionResult() {
+    doTestCannotPerform();
+  }
+
+  // PY-33843
+  public void testIteratorFreeExpressionInComprehensionIfComponent() {
+    doTest();
+  }
+
+  // PY-33843
+  public void testIteratorFreeExpressionInComprehensionResult() {
+    doTest();
+  }
+
+  // PY-33843
+  public void testExpressionInComprehensionIteratedList() {
+    doTest();
+  }
+
+  // PY-33843
+  public void testIteratorVariableDependentExpressionInComprehensionIteratedList() {
+    doTestCannotPerform();
+  }
+
   private void doTestCannotPerform() {
     boolean thrownExpectedException = false;
     try {