gh-139778: Assertion fix in code_returns_only_none

Lib/test/test_code.py

@@ -449,13 +449,21 @@ def spam5():
             if not value:
                 return None
             ...
+        def spam_while1():
+            while True:
+                pass
+        def spam_while2():
+            while True:
+                return None
         lambda1 = (lambda: None)
         for func in [
             spam1,
             spam2,
             spam3,
             spam4,
             spam5,
+            spam_while1,
+            spam_while2,
             lambda1,
         ]:
             with self.subTest(func):
@@ -474,12 +482,24 @@ def spam9():
             if value:
                 return True
             return None
+        def spam_while3():
+            i = 0
+            while True:
+                if i > 5:
+                    return i
+                else:
+                    i += 1
+        def spam_while4():
+            while True:
+                return True
         lambda2 = (lambda: True)
         for func in [
             spam6,
             spam7,
             spam8,
             spam9,
+            spam_while3,
+            spam_while4,
             lambda2,
         ]:
             with self.subTest(func):

Objects/codeobject.c

@@ -2097,10 +2097,6 @@ code_returns_only_none(PyCodeObject *co)
     int len = (int)Py_SIZE(co);
     assert(len > 0);
 
-    // The last instruction either returns or raises.  We can take advantage
-    // of that for a quick exit.
-    _Py_CODEUNIT final = _Py_GetBaseCodeUnit(co, len-1);
-
     // Look up None in co_consts.
     Py_ssize_t nconsts = PyTuple_Size(co->co_consts);
     int none_index = 0;
@@ -2113,29 +2109,23 @@ code_returns_only_none(PyCodeObject *co)
         // None wasn't there, which means there was no implicit return,
         // "return", or "return None".
 
-        // That means there must be
-        // an explicit return (non-None), or it only raises.
+        // The last instruction of a function often equals to "return X".
+        // We can take advantage of that for a quick exit.
+        _Py_CODEUNIT final = _Py_GetBaseCodeUnit(co, len-1);
         if (IS_RETURN_OPCODE(final.op.code)) {
             // It was an explicit return (non-None).
             return 0;
         }
-        // It must end with a raise then.  We still have to walk the
-        // bytecode to see if there's any explicit return (non-None).
-        assert(IS_RAISE_OPCODE(final.op.code));
-        for (int i = 0; i < len; i += _PyInstruction_GetLength(co, i)) {
-            _Py_CODEUNIT inst = _Py_GetBaseCodeUnit(co, i);
-            if (IS_RETURN_OPCODE(inst.op.code)) {
-                // We alraedy know it isn't returning None.
-                return 0;
-            }
-        }
-        // It must only raise.
+
+        // There is no need of value checking in this case.
+        none_index = -1;
     }
-    else {
-        // Walk the bytecode, looking for RETURN_VALUE.
-        for (int i = 0; i < len; i += _PyInstruction_GetLength(co, i)) {
-            _Py_CODEUNIT inst = _Py_GetBaseCodeUnit(co, i);
-            if (IS_RETURN_OPCODE(inst.op.code)) {
+
+    // Walk the bytecode, looking for RETURN_VALUE.
+    for (int i = 0; i < len; i += _PyInstruction_GetLength(co, i)) {
+        _Py_CODEUNIT inst = _Py_GetBaseCodeUnit(co, i);
+        if (IS_RETURN_OPCODE(inst.op.code)) {
+            if (none_index >= 0) {
                 assert(i != 0);
                 // Ignore it if it returns None.
                 _Py_CODEUNIT prev = _Py_GetBaseCodeUnit(co, i-1);
@@ -2145,8 +2135,8 @@ code_returns_only_none(PyCodeObject *co)
                         continue;
                     }
                 }
-                return 0;
             }
+            return 0;
         }
     }
     return 1;