[PY-26418] Fix removing whole with statement if target is unused

GitOrigin-RevId: 217ee91cebf600fa0035c2534d1b00878cf41858

python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/PyRemoveWithPartQuickFix.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2000-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.inspections.quickfix;
+
+import com.intellij.codeInspection.LocalQuickFix;
+import com.intellij.codeInspection.ProblemDescriptor;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.PyPsiBundle;
+import com.jetbrains.python.psi.*;
+import org.jetbrains.annotations.NotNull;
+
+public class PyRemoveWithPartQuickFix implements LocalQuickFix {
+  @NotNull
+  @Override
+  public String getFamilyName() {
+    return PyPsiBundle.message("QFIX.NAME.remove.target.expr");
+  }
+
+  @Override
+  public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
+    final PsiElement element = descriptor.getPsiElement();
+    final PyWithItem withItem = PsiTreeUtil.getParentOfType(element, PyWithItem.class);
+    if (withItem == null) return;
+    final PyExpression withTarget = withItem.getTarget();
+    if (withTarget != element) return;
+    final PyExpression withExpr = withItem.getExpression();
+    if (withExpr == null) return;
+    withItem.deleteChildRange(withExpr.getNextSibling(), withItem.getLastChild());
+  }
+}

python/python-psi-impl/src/com/jetbrains/python/inspections/unusedLocal/PyUnusedLocalInspectionVisitor.java

@@ -26,6 +26,7 @@ import com.jetbrains.python.inspections.quickfix.AddFieldQuickFix;
 import com.jetbrains.python.inspections.quickfix.PyRemoveExceptionTargetQuickFix;
 import com.jetbrains.python.inspections.quickfix.PyRemoveParameterQuickFix;
 import com.jetbrains.python.inspections.quickfix.PyRemoveStatementQuickFix;
+import com.jetbrains.python.inspections.quickfix.*;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.*;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
@@ -395,22 +396,34 @@ public class PyUnusedLocalInspectionVisitor extends PyInspectionVisitor {
           if (myIgnoreVariablesStartingWithUnderscore && element.getText().startsWith(PyNames.UNDERSCORE)) continue;
           if (myIgnoreTupleUnpacking && isTupleUnpacking(element)) continue;
 
+          final String warningMsg = PyPsiBundle.message("INSP.unused.locals.local.variable.isnot.used", name);
+
           final PyForStatement forStatement = PyForStatementNavigator.getPyForStatementByIterable(element);
           if (forStatement != null) {
             if (!myIgnoreRangeIterationVariables || !isRangeIteration(forStatement)) {
-              registerProblem(element, PyPsiBundle.message("INSP.unused.locals.local.variable.isnot.used", name),
-                              ProblemHighlightType.LIKE_UNUSED_SYMBOL, null, new ReplaceWithWildCard());
+              registerWarning(element, warningMsg, new ReplaceWithWildCard());
             }
             continue;
           }
 
           final PyExceptPart exceptPart = PyExceptPartNavigator.getPyExceptPartByTarget(element);
           if (exceptPart != null) {
-            registerWarning(element, PyPsiBundle.message("INSP.unused.locals.local.variable.isnot.used", name), new PyRemoveExceptionTargetQuickFix());
+            registerWarning(element, warningMsg, new PyRemoveExceptionTargetQuickFix());
             continue;
           }
 
-          registerWarning(element, PyPsiBundle.message("INSP.unused.locals.local.variable.isnot.used", name), new PyRemoveStatementQuickFix());
+          final PyWithItem withItem = PsiTreeUtil.getParentOfType(element, PyWithItem.class);
+          if (withItem != null && PsiTreeUtil.isAncestor(withItem.getTarget(), element, false)) {
+            if (withItem.getTarget() == element) {
+              registerWarning(element, warningMsg, new PyRemoveWithPartQuickFix());
+            }
+            else {
+              registerWarning(element, warningMsg, new ReplaceWithWildCard());
+            }
+            continue;
+          }
+
+          registerWarning(element, warningMsg, new PyRemoveStatementQuickFix());
         }
       }
     }

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/twoWithItemsFirstUnused.py

@@ -0,0 +1,3 @@
+def func():
+    with open('file1.txt') as <caret>unused, open('file2.txt') as used:
+        print(used)

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/twoWithItemsFirstUnused_after.py

@@ -0,0 +1,3 @@
+def func():
+    with open('file1.txt'), open('file2.txt') as used:
+        print(used)

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/twoWithItemsSecondUnused.py

@@ -0,0 +1,3 @@
+def func():
+    with open('file1.txt') as used, open('file2.txt') as <caret>unused:
+        print(used)

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/twoWithItemsSecondUnused_after.py

@@ -0,0 +1,3 @@
+def func():
+    with open('file1.txt') as used, open('file2.txt'):
+        print(used)

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/withOneTarget.py

@@ -0,0 +1,3 @@
+def main():
+    with open('file.txt') as <caret>file:
+        print(42)

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/withOneTarget_after.py

@@ -0,0 +1,3 @@
+def main():
+    with open('file.txt'):
+        print(42)

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/withTwoTargets.py

@@ -0,0 +1,11 @@
+import contextlib
+
+
+@contextlib.contextmanager
+def manager():
+    yield 1, 2
+
+
+def func():
+    with manager() as (<caret>a, b):
+        print(42)

python/testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/withTwoTargets_after.py

@@ -0,0 +1,11 @@
+import contextlib
+
+
+@contextlib.contextmanager
+def manager():
+    yield 1, 2
+
+
+def func():
+    with manager() as (_, b):
+        print(42)

python/testSrc/com/jetbrains/python/quickFixes/PyRemoveUnusedLocalQuickFixTest.java

@@ -21,8 +21,8 @@ import com.jetbrains.python.PyQuickFixTestCase;
 import com.jetbrains.python.inspections.unusedLocal.PyUnusedLocalInspection;
 import com.jetbrains.python.psi.LanguageLevel;
 
-@TestDataPath("$CONTENT_ROOT/../testData//quickFixes/PyRemoveTargetExpressionQuickFixTest/")
-public class PyRemoveExceptionTargetQuickFixTest extends PyQuickFixTestCase {
+@TestDataPath("$CONTENT_ROOT/../testData/quickFixes/PyRemoveUnusedLocalQuickFixTest/")
+public class PyRemoveUnusedLocalQuickFixTest extends PyQuickFixTestCase {
   // PY-20893
   public void testExcept() {
     doQuickFixTest(PyUnusedLocalInspection.class, PyPsiBundle.message("QFIX.NAME.remove.target.expr"));
@@ -33,4 +33,29 @@ public class PyRemoveExceptionTargetQuickFixTest extends PyQuickFixTestCase {
     runWithLanguageLevel(LanguageLevel.PYTHON27,
                          () -> doQuickFixTest(PyUnusedLocalInspection.class, PyPsiBundle.message("QFIX.NAME.remove.target.expr")));
   }
+
+  // PY-26418
+  public void testWithOneTarget() {
+    doQuickFixTest(PyUnusedLocalInspection.class, PyPsiBundle.message("QFIX.NAME.remove.target.expr"));
+  }
+
+  // PY-26418
+  public void testWithTwoTargets() {
+    doQuickFixTest(PyUnusedLocalInspection.class, PyPsiBundle.message("INSP.unused.locals.replace.with.wildcard"));
+  }
+
+  // PY-26418
+  // TODO: PY-43505
+  public void _testTwoWithItemsFirstUnused() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), () -> {
+      doQuickFixTest(PyUnusedLocalInspection.class, PyPsiBundle.message("QFIX.NAME.remove.target.expr"));
+    });
+  }
+
+  // PY-26418
+  public void testTwoWithItemsSecondUnused() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), () -> {
+      doQuickFixTest(PyUnusedLocalInspection.class, PyPsiBundle.message("QFIX.NAME.remove.target.expr"));
+    });
+  }
 }