[python] Introduced PyRaiseInstruction to be able to handle assertions while finding out function's exit points

GitOrigin-RevId: 88cb616037af7f3b833351b29b123c34fc177eec

python/python-psi-impl/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java

@@ -560,7 +560,7 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
 
   @Override
   public void visitPyRaiseStatement(final @NotNull PyRaiseStatement node) {
-    myBuilder.startNode(node);
+    myBuilder.addNodeAndCheckPending(new PyRaiseInstruction(myBuilder, node));
     final PyExpression[] expressions = node.getExpressions();
     for (PyExpression expression : expressions) {
       expression.accept(this);
@@ -834,7 +834,12 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     TransparentInstruction trueNode = addTransparentInstruction();
     TransparentInstruction falseNode = addTransparentInstruction();
     visitCondition(args[0], trueNode, falseNode);
-    myBuilder.addPendingEdge(null, falseNode);
+
+    PyRaiseInstruction raiseInstruction = new PyRaiseInstruction(myBuilder, null);
+    myBuilder.instructions.add(raiseInstruction);
+    myBuilder.addEdge(falseNode, raiseInstruction);
+
+    myBuilder.addPendingEdge(null, raiseInstruction);
     myBuilder.prevInstruction = trueNode;
   }
 

python/python-psi-impl/src/com/jetbrains/python/codeInsight/controlflow/PyRaiseInstruction.kt

@@ -0,0 +1,7 @@
+package com.jetbrains.python.codeInsight.controlflow
+
+import com.intellij.codeInsight.controlflow.ControlFlowBuilder
+import com.intellij.codeInsight.controlflow.impl.InstructionImpl
+import com.jetbrains.python.psi.PyRaiseStatement
+
+class PyRaiseInstruction(builder: ControlFlowBuilder, raiseStatement: PyRaiseStatement?): InstructionImpl(builder, raiseStatement)

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

@@ -24,6 +24,7 @@ import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.codeInsight.controlflow.CallInstruction;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
+import com.jetbrains.python.codeInsight.controlflow.PyRaiseInstruction;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider;
@@ -37,13 +38,11 @@ import com.jetbrains.python.psi.stubs.PyFunctionStub;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
 import com.jetbrains.python.psi.types.*;
 import com.jetbrains.python.sdk.PythonSdkUtil;
-import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import javax.swing.*;
 import java.util.*;
-import java.util.stream.Stream;
 
 import static com.intellij.openapi.util.text.StringUtil.notNullize;
 import static com.intellij.util.containers.ContainerUtil.*;
@@ -389,13 +388,13 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
       if (instruction instanceof CallInstruction ci && ci.isNoReturnCall(context)) {
         return ControlFlowUtil.Operation.CONTINUE;
       }
+      if (instruction instanceof PyRaiseInstruction) {
+        return ControlFlowUtil.Operation.CONTINUE;
+      }
       final PsiElement element = instruction.getElement();
       if (!(element instanceof PyStatement statement)) {
         return ControlFlowUtil.Operation.NEXT;
       }
-      if (element instanceof PyRaiseStatement || element instanceof PyAssertStatement) {
-        return ControlFlowUtil.Operation.CONTINUE;
-      }
       returnPoints.add(statement);
       return ControlFlowUtil.Operation.CONTINUE;
     });

python/testData/codeInsight/controlflow/Assert.txt

@@ -15,5 +15,6 @@
 14(15) element: null. Condition: isinstance(s, basestring):false
 15(18) ASSERTTYPE ACCESS: s
 16(17) element: null. Condition: isinstance(s, basestring):true
-17(18) ASSERTTYPE ACCESS: s
-18() element: null
+17(19) ASSERTTYPE ACCESS: s
+18(19) element: null
+19() element: null

python/testData/codeInsight/controlflow/WithAssert.txt

@@ -6,8 +6,9 @@
 5(6) READ ACCESS: f
 6(7) element: PyCallExpression: f
 7(8,9) READ ACCESS: True
-8(12) element: null. Condition: True:false
-9(10) element: null. Condition: True:true
-10(11) element: PyPrintStatement
+8(10) element: null. Condition: True:false
+9(11) element: null. Condition: True:true
+10(13) element: null
 11(12) element: PyPrintStatement
-12() element: null
+12(13) element: PyPrintStatement
+13() element: null

python/testSrc/com/jetbrains/python/Py3TypeTest.java

@@ -3233,6 +3233,14 @@ public class Py3TypeTest extends PyTestCase {
       """);
   }
 
+  public void testFunctionReturnsNone() {
+    doTest("(p: Any) -> None", """
+      def foo(p):
+          assert p
+      expr = foo
+      """);
+  }
+
   private void doTest(final String expectedType, final String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     final PyExpression expr = myFixture.findElementByText("expr", PyExpression.class);