gh-98442: Fix location info of return statement

Author: iritkatriel

Lib/test/test_compile.py

@@ -1249,6 +1249,19 @@ def test_multiline_assert(self):
         self.assertOpcodeSourcePositionIs(compiled_code, 'RAISE_VARARGS',
             line=1, end_line=3, column=0, end_column=30, occurrence=1)
 
+    def test_return_inside_with(self):
+        snippet = """\
+def f(x):
+    with x:
+        return 42
+"""
+        compiled_code, _ = self.check_positions_against_ast(snippet)
+        g = {}
+        exec(compiled_code, g)
+        compiled_code = g['f'].__code__
+        self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_VALUE',
+            line=3, end_line=3, column=8, end_column=17)
+
     def test_very_long_line_end_offset(self):
         # Make sure we get the correct column offset for offsets
         # too large to store in a byte.

Python/compile.c

@@ -3194,44 +3194,41 @@ compiler_while(struct compiler *c, stmt_ty s)
 static int
 compiler_return(struct compiler *c, stmt_ty s)
 {
-    location loc = LOC(s);
     int preserve_tos = ((s->v.Return.value != NULL) &&
                         (s->v.Return.value->kind != Constant_kind));
-    if (c->u->u_ste->ste_type != FunctionBlock)
-        return compiler_error(c, loc, "'return' outside function");
+    if (c->u->u_ste->ste_type != FunctionBlock) {
+        return compiler_error(c, LOC(s), "'return' outside function");
+    }
     if (s->v.Return.value != NULL &&
         c->u->u_ste->ste_coroutine && c->u->u_ste->ste_generator)
     {
-            return compiler_error(
-                c, loc, "'return' with value in async generator");
+        return compiler_error(
+            c, LOC(s), "'return' with value in async generator");
     }
 
     if (preserve_tos) {
         VISIT(c, expr, s->v.Return.value);
     } else {
         /* Emit instruction with line number for return value */
         if (s->v.Return.value != NULL) {
-            SET_LOC(c, s->v.Return.value);
-            loc = LOC(s->v.Return.value);
-            ADDOP(c, loc, NOP);
+            ADDOP(c, LOC(s->v.Return.value), NOP);
         }
     }
     if (s->v.Return.value == NULL || s->v.Return.value->lineno != s->lineno) {
-        SET_LOC(c, s);
-        loc = LOC(s);
-        ADDOP(c, loc, NOP);
+        ADDOP(c, LOC(s), NOP);
     }
 
-    if (!compiler_unwind_fblock_stack(c, &loc, preserve_tos, NULL))
+    location loc = LOC(s);
+    if (!compiler_unwind_fblock_stack(c, &loc, preserve_tos, NULL)) {
         return 0;
+    }
     if (s->v.Return.value == NULL) {
-        ADDOP_LOAD_CONST(c, loc, Py_None);
+        ADDOP_LOAD_CONST(c, LOC(s), Py_None);
     }
     else if (!preserve_tos) {
-        ADDOP_LOAD_CONST(c, loc, s->v.Return.value->v.Constant.value);
+        ADDOP_LOAD_CONST(c, LOC(s), s->v.Return.value->v.Constant.value);
     }
-    ADDOP(c, loc, RETURN_VALUE);
-
+    ADDOP(c, LOC(s), RETURN_VALUE);
     return 1;
 }