PY-80493 Inconsistent return inspection adding "return None" after every line in context manager

Merge-request: IJ-MR-160763
Merged-by: Aleksandr Govenko <aleksandr.govenko@jetbrains.com>

GitOrigin-RevId: b576e30d433d39e426c8854d3ad485594ac9e958

python/python-psi-impl/src/com/jetbrains/python/inspections/PyInconsistentReturnsInspection.kt

@@ -14,6 +14,8 @@ import com.jetbrains.python.psi.PyReturnStatement
  * or none of them should. If any return statement returns an expression, any return statements where no value
  * is returned should explicitly state this as return None, and an explicit return statement should be present
  * at the end of the function (if reachable).
+ * 
+ * @see PyMakeReturnsExplicitFix for tests
  */
 class PyInconsistentReturnsInspection : PyInspection() {
   override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean, session: LocalInspectionToolSession): PsiElementVisitor {

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java

@@ -405,7 +405,10 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
         if (instruction instanceof PyRaiseInstruction) {
           return ControlFlowUtil.Operation.CONTINUE;
         }
-        if (instruction instanceof PyWithContextExitInstruction withExit && !withExit.isSuppressingExceptions(context)) {
+        if (instruction instanceof PyWithContextExitInstruction withExit) {
+          if (collectImplicitReturn && withExit.isSuppressingExceptions(context)) {
+            returnPoints.add(PsiTreeUtil.getParentOfType(withExit.getElement(), PyWithStatement.class));
+          }
           return ControlFlowUtil.Operation.CONTINUE;
         }
         final PsiElement element = instruction.getElement();

python/testData/quickFixes/PyMakeReturnsExplicitFixTest/addReturnsFromAnnotation.py

@@ -1,5 +0,0 @@
-def f(x) -> <warning descr="Function returning 'int | None' has implicit return">int | None<caret></warning>:
-    if x == 1:
-        return 42
-    elif x == 2:
-        <warning descr="Function returning 'int | None' has implicit return">return</warning>

python/testData/quickFixes/PyMakeReturnsExplicitFixTest/addReturnsFromAnnotation_after.py

@@ -1,6 +0,0 @@
-def f(x) -> int | None:
-    if x == 1:
-        return 42
-    elif x == 2:
-        return None
-    return None

python/testData/quickFixes/PyMakeReturnsExplicitFixTest/contextManagerNotSuppressingException.py

@@ -0,0 +1,11 @@
+class NotSuppressingContext:
+    def __enter__(self):
+        ...
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool | None:
+        ...
+
+def foo():
+    with NotSuppressingContext() as st:
+        foo()
+        <weak_warning descr="Explicit return statement expected">if bool():
+            return 1<caret></weak_warning>

python/testData/quickFixes/PyMakeReturnsExplicitFixTest/contextManagerNotSuppressingException_after.py

@@ -0,0 +1,12 @@
+class NotSuppressingContext:
+    def __enter__(self):
+        ...
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool | None:
+        ...
+
+def foo():
+    with NotSuppressingContext() as st:
+        foo()
+        if bool():
+            return 1
+        return None

python/testData/quickFixes/PyMakeReturnsExplicitFixTest/contextManagerSuppressingException.py

@@ -0,0 +1,10 @@
+class SuppressingContext:
+    def __enter__(self):
+        ...
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        ...
+
+def foo():
+    <weak_warning descr="Explicit return statement expected">with SuppressingContext() as st:
+        foo()
+        return 1<caret></weak_warning>

python/testData/quickFixes/PyMakeReturnsExplicitFixTest/contextManagerSuppressingException_after.py

@@ -0,0 +1,11 @@
+class SuppressingContext:
+    def __enter__(self):
+        ...
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        ...
+
+def foo():
+    with SuppressingContext() as st:
+        foo()
+        return 1
+    return None

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

@@ -13,6 +13,16 @@ public class PyMakeReturnsExplicitFixTest extends PyQuickFixTestCase {
     doQuickFixTest(PyInconsistentReturnsInspection.class, PyPsiBundle.message("QFIX.NAME.make.return.stmts.explicit"));
   }
 
+  // PY-80493
+  public void testContextManagerSuppressingException() {
+    doQuickFixTest(PyInconsistentReturnsInspection.class, PyPsiBundle.message("QFIX.NAME.make.return.stmts.explicit"));
+  }
+
+  // PY-80493
+  public void testContextManagerNotSuppressingException() {
+    doQuickFixTest(PyInconsistentReturnsInspection.class, PyPsiBundle.message("QFIX.NAME.make.return.stmts.explicit"));
+  }
+
   @Override
   protected void doQuickFixTest(final Class inspectionClass, final String hint) {
     final String testFileName = getTestName(true);