From 92aa9db8c47d176d797a16150568cd7820a19daf Mon Sep 17 00:00:00 2001
From: "Jeong, YunWon" <jeong@youknowone.org>
Date: Thu, 29 Jan 2026 23:14:33 +0900
Subject: [PATCH] Align GetAwaitable to Python 3.14.2

---
 Lib/test/test_asyncio/test_locks.py           | 34 ++++----
 Lib/test/test_asyncio/test_pep492.py          |  7 +-
 crates/codegen/src/compile.rs                 | 14 +--
 ...pile__tests__nested_double_async_with.snap |  7 +-
 .../compiler-core/src/bytecode/instruction.rs |  8 +-
 crates/vm/src/coroutine.rs                    | 46 ++++++++++
 crates/vm/src/frame.rs                        | 85 +++++++++----------
 7 files changed, 121 insertions(+), 80 deletions(-)

diff --git a/Lib/test/test_asyncio/test_locks.py b/Lib/test/test_asyncio/test_locks.py
index a0884bffe6b..e025d2990a3 100644
--- a/Lib/test/test_asyncio/test_locks.py
+++ b/Lib/test/test_asyncio/test_locks.py
@@ -14,24 +14,24 @@
     r'(, value:\d)?'
     r'(, waiters:\d+)?'
     r'(, waiters:\d+\/\d+)?' # barrier
-    r')\]>\Z'
+    r')\]>\z'
 )
 RGX_REPR = re.compile(STR_RGX_REPR)
 
 
 def tearDownModule():
-    asyncio.set_event_loop_policy(None)
+    asyncio.events._set_event_loop_policy(None)
 
 
 class LockTests(unittest.IsolatedAsyncioTestCase):
 
     async def test_repr(self):
         lock = asyncio.Lock()
-        self.assertTrue(repr(lock).endswith('[unlocked]>'))
+        self.assertEndsWith(repr(lock), '[unlocked]>')
         self.assertTrue(RGX_REPR.match(repr(lock)))
 
         await lock.acquire()
-        self.assertTrue(repr(lock).endswith('[locked]>'))
+        self.assertEndsWith(repr(lock), '[locked]>')
         self.assertTrue(RGX_REPR.match(repr(lock)))
 
     async def test_lock(self):
@@ -39,7 +39,7 @@ async def test_lock(self):
 
         with self.assertRaisesRegex(
             TypeError,
-            "object Lock can't be used in 'await' expression"
+            "'Lock' object can't be awaited"
         ):
             await lock
 
@@ -77,7 +77,7 @@ async def test_lock_by_with_statement(self):
             self.assertFalse(lock.locked())
             with self.assertRaisesRegex(
                 TypeError,
-                r"object \w+ can't be used in 'await' expression"
+                r"'\w+' object can't be awaited"
             ):
                 with await lock:
                     pass
@@ -286,12 +286,12 @@ class EventTests(unittest.IsolatedAsyncioTestCase):
 
     def test_repr(self):
         ev = asyncio.Event()
-        self.assertTrue(repr(ev).endswith('[unset]>'))
+        self.assertEndsWith(repr(ev), '[unset]>')
         match = RGX_REPR.match(repr(ev))
         self.assertEqual(match.group('extras'), 'unset')
 
         ev.set()
-        self.assertTrue(repr(ev).endswith('[set]>'))
+        self.assertEndsWith(repr(ev), '[set]>')
         self.assertTrue(RGX_REPR.match(repr(ev)))
 
         ev._waiters.append(mock.Mock())
@@ -916,11 +916,11 @@ def test_initial_value_zero(self):
 
     async def test_repr(self):
         sem = asyncio.Semaphore()
-        self.assertTrue(repr(sem).endswith('[unlocked, value:1]>'))
+        self.assertEndsWith(repr(sem), '[unlocked, value:1]>')
         self.assertTrue(RGX_REPR.match(repr(sem)))
 
         await sem.acquire()
-        self.assertTrue(repr(sem).endswith('[locked]>'))
+        self.assertEndsWith(repr(sem), '[locked]>')
         self.assertTrue('waiters' not in repr(sem))
         self.assertTrue(RGX_REPR.match(repr(sem)))
 
@@ -941,7 +941,7 @@ async def test_semaphore(self):
 
         with self.assertRaisesRegex(
             TypeError,
-            "object Semaphore can't be used in 'await' expression",
+            "'Semaphore' object can't be awaited",
         ):
             await sem
 
@@ -1194,14 +1194,14 @@ async def c3(result):
         self.assertEqual([2, 3], result)
 
     async def test_acquire_fifo_order_4(self):
-        # Test that a successfule `acquire()` will wake up multiple Tasks
+        # Test that a successful `acquire()` will wake up multiple Tasks
         # that were waiting in the Semaphore queue due to FIFO rules.
         sem = asyncio.Semaphore(0)
         result = []
         count = 0
 
         async def c1(result):
-            # First task immediatlly waits for semaphore.  It will be awoken by c2.
+            # First task immediately waits for semaphore.  It will be awoken by c2.
             self.assertEqual(sem._value, 0)
             await sem.acquire()
             # We should have woken up all waiting tasks now.
@@ -1270,7 +1270,7 @@ async def test_barrier(self):
         self.assertIn("filling", repr(barrier))
         with self.assertRaisesRegex(
             TypeError,
-            "object Barrier can't be used in 'await' expression",
+            "'Barrier' object can't be awaited",
         ):
             await barrier
 
@@ -1475,13 +1475,13 @@ async def coro():
             # first time waiting
             await barrier.wait()
 
-            # after wainting once for all tasks
+            # after waiting once for all tasks
             if rewait_n > 0:
                 rewait_n -= 1
                 # wait again only for rewait tasks
                 await barrier.wait()
             else:
-                # wait for end of draining state`
+                # wait for end of draining state
                 await barrier_nowaiting.wait()
                 # wait for other waiting tasks
                 await barrier.wait()
@@ -1780,7 +1780,7 @@ async def coro():
         self.assertEqual(barrier.n_waiting, 0)
 
     async def test_abort_barrier_when_exception_then_resetting(self):
-        # test from threading.Barrier: see `lock_tests.test_abort_and_reset``
+        # test from threading.Barrier: see `lock_tests.test_abort_and_reset`
         barrier1 = asyncio.Barrier(self.N)
         barrier2 = asyncio.Barrier(self.N)
         results1 = []
diff --git a/Lib/test/test_asyncio/test_pep492.py b/Lib/test/test_asyncio/test_pep492.py
index dc25a46985e..95bdc283df0 100644
--- a/Lib/test/test_asyncio/test_pep492.py
+++ b/Lib/test/test_asyncio/test_pep492.py
@@ -11,7 +11,7 @@
 
 
 def tearDownModule():
-    asyncio.set_event_loop_policy(None)
+    asyncio.events._set_event_loop_policy(None)
 
 
 # Test that asyncio.iscoroutine() uses collections.abc.Coroutine
@@ -77,7 +77,7 @@ async def test(lock):
             self.assertFalse(lock.locked())
             with self.assertRaisesRegex(
                 TypeError,
-                "can't be used in 'await' expression"
+                "can't be awaited"
             ):
                 with await lock:
                     pass
@@ -124,9 +124,10 @@ def foo(): yield
 
         self.assertFalse(asyncio.iscoroutine(foo()))
 
-
     def test_iscoroutinefunction(self):
         async def foo(): pass
+        # TODO: RUSTPYTHON; no DeprecationWarning for iscoroutinefunction
+        # with self.assertWarns(DeprecationWarning):
         self.assertTrue(asyncio.iscoroutinefunction(foo))
 
     def test_async_def_coroutines(self):
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
index b961be42517..12ba24253cf 100644
--- a/crates/codegen/src/compile.rs
+++ b/crates/codegen/src/compile.rs
@@ -1465,7 +1465,7 @@ impl Compiler {
 
                 // For async with, await the result
                 if matches!(info.fb_type, FBlockType::AsyncWith) {
-                    emit!(self, Instruction::GetAwaitable);
+                    emit!(self, Instruction::GetAwaitable { arg: 2 });
                     self.emit_load_const(ConstantData::None);
                     self.compile_yield_from_sequence(true)?;
                 }
@@ -4773,7 +4773,7 @@ impl Compiler {
             // bound_aenter is already bound, call with NULL self_or_null
             emit!(self, Instruction::PushNull); // [bound_aexit, bound_aenter, NULL]
             emit!(self, Instruction::Call { nargs: 0 }); // [bound_aexit, awaitable]
-            emit!(self, Instruction::GetAwaitable);
+            emit!(self, Instruction::GetAwaitable { arg: 1 });
             self.emit_load_const(ConstantData::None);
             self.compile_yield_from_sequence(true)?;
         } else {
@@ -4854,7 +4854,7 @@ impl Compiler {
         self.emit_load_const(ConstantData::None);
         emit!(self, Instruction::Call { nargs: 3 });
         if is_async {
-            emit!(self, Instruction::GetAwaitable);
+            emit!(self, Instruction::GetAwaitable { arg: 2 });
             self.emit_load_const(ConstantData::None);
             self.compile_yield_from_sequence(true)?;
         }
@@ -4899,7 +4899,7 @@ impl Compiler {
         emit!(self, Instruction::WithExceptStart);
 
         if is_async {
-            emit!(self, Instruction::GetAwaitable);
+            emit!(self, Instruction::GetAwaitable { arg: 2 });
             self.emit_load_const(ConstantData::None);
             self.compile_yield_from_sequence(true)?;
         }
@@ -6741,7 +6741,7 @@ impl Compiler {
                     return Err(self.error(CodegenErrorType::InvalidAwait));
                 }
                 self.compile_expression(value)?;
-                emit!(self, Instruction::GetAwaitable);
+                emit!(self, Instruction::GetAwaitable { arg: 0 });
                 self.emit_load_const(ConstantData::None);
                 self.compile_yield_from_sequence(true)?;
             }
@@ -7516,7 +7516,7 @@ impl Compiler {
         // Call just created <listcomp> function:
         emit!(self, Instruction::Call { nargs: 1 });
         if is_async_list_set_dict_comprehension {
-            emit!(self, Instruction::GetAwaitable);
+            emit!(self, Instruction::GetAwaitable { arg: 0 });
             self.emit_load_const(ConstantData::None);
             self.compile_yield_from_sequence(true)?;
         }
@@ -7991,7 +7991,7 @@ impl Compiler {
                     emit!(self, Instruction::Call { nargs: 3 });
 
                     if is_async {
-                        emit!(self, Instruction::GetAwaitable);
+                        emit!(self, Instruction::GetAwaitable { arg: 2 });
                         self.emit_load_const(ConstantData::None);
                         self.compile_yield_from_sequence(true)?;
                     }
diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap
index 4f67bd10805..28a84bc1e86 100644
--- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap
+++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap
@@ -1,6 +1,5 @@
 ---
 source: crates/codegen/src/compile.rs
-assertion_line: 8782
 expression: "compile_exec(\"\\\nasync def test():\n    for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')):\n        with self.subTest(type=type(stop_exc)):\n            try:\n                async with egg():\n                    raise stop_exc\n            except Exception as ex:\n                self.assertIs(ex, stop_exc)\n            else:\n                self.fail(f'{stop_exc} was suppressed')\n\")"
 ---
   1           0 RESUME               (0)
@@ -47,7 +46,7 @@ expression: "compile_exec(\"\\\nasync def test():\n    for stop_exc in (StopIter
                  36 LOAD_SPECIAL         (__aenter__)
                  37 PUSH_NULL
                  38 CALL                 (0)
-                 39 GET_AWAITABLE
+                 39 GET_AWAITABLE        (1)
                  40 LOAD_CONST           (None)
             >>   41 SEND                 (46)
                  42 YIELD_VALUE          (1)
@@ -65,7 +64,7 @@ expression: "compile_exec(\"\\\nasync def test():\n    for stop_exc in (StopIter
                  52 LOAD_CONST           (None)
                  53 LOAD_CONST           (None)
                  54 CALL                 (3)
-                 55 GET_AWAITABLE
+                 55 GET_AWAITABLE        (2)
                  56 LOAD_CONST           (None)
             >>   57 SEND                 (62)
                  58 YIELD_VALUE          (1)
@@ -77,7 +76,7 @@ expression: "compile_exec(\"\\\nasync def test():\n    for stop_exc in (StopIter
                  64 JUMP_FORWARD         (86)
                  65 PUSH_EXC_INFO
                  66 WITH_EXCEPT_START
-                 67 GET_AWAITABLE
+                 67 GET_AWAITABLE        (2)
                  68 LOAD_CONST           (None)
             >>   69 SEND                 (74)
                  70 YIELD_VALUE          (1)
diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs
index 4bf15cf5772..0458f5e4baf 100644
--- a/crates/compiler-core/src/bytecode/instruction.rs
+++ b/crates/compiler-core/src/bytecode/instruction.rs
@@ -142,7 +142,9 @@ pub enum Instruction {
     ForIter {
         target: Arg<Label>,
     } = 70,
-    GetAwaitable = 71, // TODO: Make this instruction to hold an oparg
+    GetAwaitable {
+        arg: Arg<u32>,
+    } = 71,
     ImportFrom {
         idx: Arg<NameIdx>,
     } = 72,
@@ -582,7 +584,7 @@ impl InstructionMetadata for Instruction {
             }
             Self::PopExcept => -1,
             Self::PopIter => -1,
-            Self::GetAwaitable => 0,
+            Self::GetAwaitable { .. } => 0,
             Self::GetAIter => 0,
             Self::GetANext => 1,
             Self::EndAsyncFor => -2,  // pops (awaitable, exc) from stack
@@ -828,7 +830,7 @@ impl InstructionMetadata for Instruction {
             Self::FormatWithSpec => w!(FORMAT_WITH_SPEC),
             Self::GetAIter => w!(GET_AITER),
             Self::GetANext => w!(GET_ANEXT),
-            Self::GetAwaitable => w!(GET_AWAITABLE),
+            Self::GetAwaitable { arg } => w!(GET_AWAITABLE, arg),
             Self::Reserved => w!(RESERVED),
             Self::GetIter => w!(GET_ITER),
             Self::GetLen => w!(GET_LEN),
diff --git a/crates/vm/src/coroutine.rs b/crates/vm/src/coroutine.rs
index 236388ecf8b..67325283f3a 100644
--- a/crates/vm/src/coroutine.rs
+++ b/crates/vm/src/coroutine.rs
@@ -246,6 +246,52 @@ pub fn is_gen_exit(exc: &Py<PyBaseException>, vm: &VirtualMachine) -> bool {
     exc.fast_isinstance(vm.ctx.exceptions.generator_exit)
 }
 
+/// Get an awaitable iterator from an object.
+///
+/// Returns the object itself if it's a coroutine or iterable coroutine (generator with
+/// CO_ITERABLE_COROUTINE flag). Otherwise calls `__await__()` and validates the result.
+pub fn get_awaitable_iter(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+    use crate::builtins::{PyCoroutine, PyGenerator};
+    use crate::protocol::PyIter;
+
+    if obj.downcastable::<PyCoroutine>()
+        || obj.downcast_ref::<PyGenerator>().is_some_and(|g| {
+            g.as_coro()
+                .frame()
+                .code
+                .flags
+                .contains(crate::bytecode::CodeFlags::ITERABLE_COROUTINE)
+        })
+    {
+        return Ok(obj);
+    }
+
+    if let Some(await_method) = vm.get_method(obj.clone(), identifier!(vm, __await__)) {
+        let result = await_method?.call((), vm)?;
+        // __await__() must NOT return a coroutine (PEP 492)
+        if result.downcastable::<PyCoroutine>()
+            || result.downcast_ref::<PyGenerator>().is_some_and(|g| {
+                g.as_coro()
+                    .frame()
+                    .code
+                    .flags
+                    .contains(crate::bytecode::CodeFlags::ITERABLE_COROUTINE)
+            })
+        {
+            return Err(vm.new_type_error("__await__() returned a coroutine".to_owned()));
+        }
+        if !PyIter::check(&result) {
+            return Err(vm.new_type_error(format!(
+                "__await__() returned non-iterator of type '{}'",
+                result.class().name()
+            )));
+        }
+        return Ok(result);
+    }
+
+    Err(vm.new_type_error(format!("'{}' object can't be awaited", obj.class().name())))
+}
+
 /// Emit DeprecationWarning for the deprecated 3-argument throw() signature.
 pub fn warn_deprecated_throw_signature(
     exc_val: &OptionalArg,
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
index b7df081a642..844594f7859 100644
--- a/crates/vm/src/frame.rs
+++ b/crates/vm/src/frame.rs
@@ -1025,54 +1025,47 @@ impl ExecutingFrame<'_> {
                 debug_assert_eq!(orig_stack_len + 1, self.state.stack.len());
                 Ok(None)
             }
-            Instruction::GetAwaitable => {
-                use crate::protocol::PyIter;
-
-                let awaited_obj = self.pop_value();
-                let awaitable = if let Some(coro) = awaited_obj.downcast_ref::<PyCoroutine>() {
-                    // _PyGen_yf() check - detect if coroutine is already being awaited elsewhere
-                    if coro.as_coro().frame().yield_from_target().is_some() {
-                        return Err(
-                            vm.new_runtime_error("coroutine is being awaited already".to_owned())
-                        );
-                    }
-                    awaited_obj
-                } else if awaited_obj
-                    .downcast_ref::<PyGenerator>()
-                    .is_some_and(|generator| {
-                        generator
-                            .as_coro()
-                            .frame()
-                            .code
-                            .flags
-                            .contains(bytecode::CodeFlags::ITERABLE_COROUTINE)
-                    })
-                {
-                    // Generator with CO_ITERABLE_COROUTINE flag can be awaited
-                    // (e.g., generators decorated with @types.coroutine)
-                    awaited_obj
-                } else {
-                    let await_method = vm.get_method_or_type_error(
-                        awaited_obj.clone(),
-                        identifier!(vm, __await__),
-                        || {
-                            format!(
-                                "object {} can't be used in 'await' expression",
-                                awaited_obj.class().name(),
-                            )
-                        },
-                    )?;
-                    let result = await_method.call((), vm)?;
-                    // Check that __await__ returned an iterator
-                    if !PyIter::check(&result) {
-                        return Err(vm.new_type_error(format!(
-                            "__await__() returned non-iterator of type '{}'",
-                            result.class().name()
-                        )));
+            Instruction::GetAwaitable { arg: oparg } => {
+                let iterable = self.pop_value();
+
+                let iter = match crate::coroutine::get_awaitable_iter(iterable.clone(), vm) {
+                    Ok(iter) => iter,
+                    Err(e) => {
+                        // _PyEval_FormatAwaitableError: override error for async with
+                        // when the type doesn't have __await__
+                        let oparg_val = oparg.get(arg);
+                        if vm
+                            .get_method(iterable.clone(), identifier!(vm, __await__))
+                            .is_none()
+                        {
+                            if oparg_val == 1 {
+                                return Err(vm.new_type_error(format!(
+                                    "'async with' received an object from __aenter__ \
+                                     that does not implement __await__: {}",
+                                    iterable.class().name()
+                                )));
+                            } else if oparg_val == 2 {
+                                return Err(vm.new_type_error(format!(
+                                    "'async with' received an object from __aexit__ \
+                                     that does not implement __await__: {}",
+                                    iterable.class().name()
+                                )));
+                            }
+                        }
+                        return Err(e);
                     }
-                    result
                 };
-                self.push_value(awaitable);
+
+                // Check if coroutine is already being awaited
+                if let Some(coro) = iter.downcast_ref::<PyCoroutine>()
+                    && coro.as_coro().frame().yield_from_target().is_some()
+                {
+                    return Err(
+                        vm.new_runtime_error("coroutine is being awaited already".to_owned())
+                    );
+                }
+
+                self.push_value(iter);
                 Ok(None)
             }
             Instruction::GetIter => {