PyAnextAwaitable (#6427)

.cspell.dict/python-more.txt

@@ -4,6 +4,7 @@ aenter
 aexit
 aiter
 anext
+anextawaitable
 appendleft
 argcount
 arrayiterator

Lib/test/test_asyncgen.py

@@ -375,8 +375,6 @@ class AsyncGenTest(unittest.TestCase):
 
         self.compare_generators(sync_gen_wrapper(), async_gen_wrapper())
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_async_gen_api_01(self):
         async def gen():
             yield 123
@@ -467,16 +465,12 @@ class AsyncGenAsyncioTest(unittest.TestCase):
         result = self.loop.run_until_complete(test_throw())
         self.assertEqual(result, "completed")
 
-    # TODO: RUSTPYTHON, NameError: name 'anext' is not defined
-    @unittest.expectedFailure
     def test_async_generator_anext(self):
         async def agen():
             yield 1
             yield 2
         self.check_async_iterator_anext(agen)
 
-    # TODO: RUSTPYTHON, NameError: name 'anext' is not defined
-    @unittest.expectedFailure
     def test_python_async_iterator_anext(self):
         class MyAsyncIter:
             """Asynchronously yield 1, then 2."""
@@ -492,8 +486,6 @@ class AsyncGenAsyncioTest(unittest.TestCase):
                     return self.yielded
         self.check_async_iterator_anext(MyAsyncIter)
 
-    # TODO: RUSTPYTHON, NameError: name 'anext' is not defined
-    @unittest.expectedFailure
     def test_python_async_iterator_types_coroutine_anext(self):
         import types
         class MyAsyncIterWithTypesCoro:
@@ -523,8 +515,6 @@ class AsyncGenAsyncioTest(unittest.TestCase):
         res = self.loop.run_until_complete(consume())
         self.assertEqual(res, [1, 2])
 
-    # TODO: RUSTPYTHON, NameError: name 'aiter' is not defined
-    @unittest.expectedFailure
     def test_async_gen_aiter_class(self):
         results = []
         class Gen:
@@ -549,8 +539,6 @@ class AsyncGenAsyncioTest(unittest.TestCase):
         applied_twice = aiter(applied_once)
         self.assertIs(applied_once, applied_twice)
 
-    # TODO: RUSTPYTHON, NameError: name 'anext' is not defined
-    @unittest.expectedFailure
     def test_anext_bad_args(self):
         async def gen():
             yield 1
@@ -571,7 +559,7 @@ class AsyncGenAsyncioTest(unittest.TestCase):
         with self.assertRaises(TypeError):
             self.loop.run_until_complete(call_with_kwarg())
 
-    # TODO: RUSTPYTHON, NameError: name 'anext' is not defined
+    # TODO: RUSTPYTHON, error message mismatch
     @unittest.expectedFailure
     def test_anext_bad_await(self):
         async def bad_awaitable():
@@ -642,7 +630,7 @@ class AsyncGenAsyncioTest(unittest.TestCase):
         result = self.loop.run_until_complete(do_test())
         self.assertEqual(result, "completed")
 
-    # TODO: RUSTPYTHON, NameError: name 'anext' is not defined
+    # TODO: RUSTPYTHON, anext coroutine iteration issue
     @unittest.expectedFailure
     def test_anext_iter(self):
         @types.coroutine

Lib/test/test_generators.py

@@ -117,7 +117,6 @@ class FinalizationTest(unittest.TestCase):
 
 class GeneratorTest(unittest.TestCase):
 
-    @unittest.expectedFailure # TODO: RUSTPYTHON
     def test_name(self):
         def func():
             yield 1

Lib/test/test_types.py

@@ -2052,8 +2052,6 @@ class CoroutineTests(unittest.TestCase):
         else:
             self.fail('StopIteration was expected')
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_gen(self):
         def gen_func():
             yield 1

crates/vm/src/builtins/asyncgenerator.rs

@@ -33,9 +33,9 @@ impl PyAsyncGen {
         &self.inner
     }
 
-    pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
+    pub fn new(frame: FrameRef, name: PyStrRef, qualname: PyStrRef) -> Self {
         Self {
-            inner: Coro::new(frame, name),
+            inner: Coro::new(frame, name, qualname),
             running_async: AtomicCell::new(false),
         }
     }
@@ -50,6 +50,16 @@ impl PyAsyncGen {
         self.inner.set_name(name)
     }
 
+    #[pygetset]
+    fn __qualname__(&self) -> PyStrRef {
+        self.inner.qualname()
+    }
+
+    #[pygetset(setter)]
+    fn set___qualname__(&self, qualname: PyStrRef) {
+        self.inner.set_qualname(qualname)
+    }
+
     #[pygetset]
     fn ag_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
         self.inner.frame().yield_from_target()
@@ -424,8 +434,151 @@ impl IterNext for PyAsyncGenAThrow {
     }
 }
 
+/// Awaitable wrapper for anext() builtin with default value.
+/// When StopAsyncIteration is raised, it converts it to StopIteration(default).
+#[pyclass(module = false, name = "anext_awaitable")]
+#[derive(Debug)]
+pub struct PyAnextAwaitable {
+    wrapped: PyObjectRef,
+    default_value: PyObjectRef,
+}
+
+impl PyPayload for PyAnextAwaitable {
+    #[inline]
+    fn class(ctx: &Context) -> &'static Py<PyType> {
+        ctx.types.anext_awaitable
+    }
+}
+
+#[pyclass(with(IterNext, Iterable))]
+impl PyAnextAwaitable {
+    pub fn new(wrapped: PyObjectRef, default_value: PyObjectRef) -> Self {
+        Self {
+            wrapped,
+            default_value,
+        }
+    }
+
+    #[pymethod(name = "__await__")]
+    fn r#await(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
+        zelf
+    }
+
+    /// Get the awaitable iterator from wrapped object.
+    // = anextawaitable_getiter.
+    fn get_awaitable_iter(&self, vm: &VirtualMachine) -> PyResult {
+        use crate::builtins::PyCoroutine;
+        use crate::protocol::PyIter;
+
+        let wrapped = &self.wrapped;
+
+        // If wrapped is already an async_generator_asend, it's an iterator
+        if wrapped.class().is(vm.ctx.types.async_generator_asend)
+            || wrapped.class().is(vm.ctx.types.async_generator_athrow)
+        {
+            return Ok(wrapped.clone());
+        }
+
+        // _PyCoro_GetAwaitableIter equivalent
+        let awaitable = if wrapped.class().is(vm.ctx.types.coroutine_type) {
+            // Coroutine - get __await__ later
+            wrapped.clone()
+        } else {
+            // Try to get __await__ method
+            if let Some(await_method) = vm.get_method(wrapped.clone(), identifier!(vm, __await__)) {
+                await_method?.call((), vm)?
+            } else {
+                return Err(vm.new_type_error(format!(
+                    "object {} can't be used in 'await' expression",
+                    wrapped.class().name()
+                )));
+            }
+        };
+
+        // If awaitable is a coroutine, get its __await__
+        if awaitable.class().is(vm.ctx.types.coroutine_type) {
+            let coro_await = vm.call_method(&awaitable, "__await__", ())?;
+            // Check that __await__ returned an iterator
+            if !PyIter::check(&coro_await) {
+                return Err(vm.new_type_error("__await__ returned a non-iterable"));
+            }
+            return Ok(coro_await);
+        }
+
+        // Check the result is an iterator, not a coroutine
+        if awaitable.downcast_ref::<PyCoroutine>().is_some() {
+            return Err(vm.new_type_error("__await__() returned a coroutine"));
+        }
+
+        // Check that the result is an iterator
+        if !PyIter::check(&awaitable) {
+            return Err(vm.new_type_error(format!(
+                "__await__() returned non-iterator of type '{}'",
+                awaitable.class().name()
+            )));
+        }
+
+        Ok(awaitable)
+    }
+
+    #[pymethod]
+    fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        let awaitable = self.get_awaitable_iter(vm)?;
+        let result = vm.call_method(&awaitable, "send", (val,));
+        self.handle_result(result, vm)
+    }
+
+    #[pymethod]
+    fn throw(
+        &self,
+        exc_type: PyObjectRef,
+        exc_val: OptionalArg,
+        exc_tb: OptionalArg,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        let awaitable = self.get_awaitable_iter(vm)?;
+        let result = vm.call_method(
+            &awaitable,
+            "throw",
+            (
+                exc_type,
+                exc_val.unwrap_or_none(vm),
+                exc_tb.unwrap_or_none(vm),
+            ),
+        );
+        self.handle_result(result, vm)
+    }
+
+    #[pymethod]
+    fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
+        if let Ok(awaitable) = self.get_awaitable_iter(vm) {
+            let _ = vm.call_method(&awaitable, "close", ());
+        }
+        Ok(())
+    }
+
+    /// Convert StopAsyncIteration to StopIteration(default_value)
+    fn handle_result(&self, result: PyResult, vm: &VirtualMachine) -> PyResult {
+        match result {
+            Ok(value) => Ok(value),
+            Err(exc) if exc.fast_isinstance(vm.ctx.exceptions.stop_async_iteration) => {
+                Err(vm.new_stop_iteration(Some(self.default_value.clone())))
+            }
+            Err(exc) => Err(exc),
+        }
+    }
+}
+
+impl SelfIter for PyAnextAwaitable {}
+impl IterNext for PyAnextAwaitable {
+    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+        PyIterReturn::from_pyresult(zelf.send(vm.ctx.none(), vm), vm)
+    }
+}
+
 pub fn init(ctx: &Context) {
     PyAsyncGen::extend_class(ctx, ctx.types.async_generator);
     PyAsyncGenASend::extend_class(ctx, ctx.types.async_generator_asend);
     PyAsyncGenAThrow::extend_class(ctx, ctx.types.async_generator_athrow);
+    PyAnextAwaitable::extend_class(ctx, ctx.types.anext_awaitable);
 }

crates/vm/src/builtins/coroutine.rs

@@ -29,9 +29,9 @@ impl PyCoroutine {
         &self.inner
     }
 
-    pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
+    pub fn new(frame: FrameRef, name: PyStrRef, qualname: PyStrRef) -> Self {
         Self {
-            inner: Coro::new(frame, name),
+            inner: Coro::new(frame, name, qualname),
         }
     }
 
@@ -45,6 +45,16 @@ impl PyCoroutine {
         self.inner.set_name(name)
     }
 
+    #[pygetset]
+    fn __qualname__(&self) -> PyStrRef {
+        self.inner.qualname()
+    }
+
+    #[pygetset(setter)]
+    fn set___qualname__(&self, qualname: PyStrRef) {
+        self.inner.set_qualname(qualname)
+    }
+
     #[pymethod(name = "__await__")]
     const fn r#await(zelf: PyRef<Self>) -> PyCoroutineWrapper {
         PyCoroutineWrapper { coro: zelf }
@@ -156,6 +166,11 @@ impl PyCoroutineWrapper {
     ) -> PyResult<PyIterReturn> {
         self.coro.throw(exc_type, exc_val, exc_tb, vm)
     }
+
+    #[pymethod]
+    fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
+        self.coro.close(vm)
+    }
 }
 
 impl SelfIter for PyCoroutineWrapper {}

crates/vm/src/builtins/function.rs

@@ -425,9 +425,15 @@ impl Py<PyFunction> {
         let is_gen = code.flags.contains(bytecode::CodeFlags::IS_GENERATOR);
         let is_coro = code.flags.contains(bytecode::CodeFlags::IS_COROUTINE);
         match (is_gen, is_coro) {
-            (true, false) => Ok(PyGenerator::new(frame, self.__name__()).into_pyobject(vm)),
-            (false, true) => Ok(PyCoroutine::new(frame, self.__name__()).into_pyobject(vm)),
-            (true, true) => Ok(PyAsyncGen::new(frame, self.__name__()).into_pyobject(vm)),
+            (true, false) => {
+                Ok(PyGenerator::new(frame, self.__name__(), self.__qualname__()).into_pyobject(vm))
+            }
+            (false, true) => {
+                Ok(PyCoroutine::new(frame, self.__name__(), self.__qualname__()).into_pyobject(vm))
+            }
+            (true, true) => {
+                Ok(PyAsyncGen::new(frame, self.__name__(), self.__qualname__()).into_pyobject(vm))
+            }
             (false, false) => vm.run_frame(frame),
         }
     }

crates/vm/src/builtins/generator.rs

@@ -32,9 +32,9 @@ impl PyGenerator {
         &self.inner
     }
 
-    pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
+    pub fn new(frame: FrameRef, name: PyStrRef, qualname: PyStrRef) -> Self {
         Self {
-            inner: Coro::new(frame, name),
+            inner: Coro::new(frame, name, qualname),
         }
     }
 
@@ -48,6 +48,16 @@ impl PyGenerator {
         self.inner.set_name(name)
     }
 
+    #[pygetset]
+    fn __qualname__(&self) -> PyStrRef {
+        self.inner.qualname()
+    }
+
+    #[pygetset(setter)]
+    fn set___qualname__(&self, qualname: PyStrRef) {
+        self.inner.set_qualname(qualname)
+    }
+
     #[pygetset]
     fn gi_frame(&self, _vm: &VirtualMachine) -> FrameRef {
         self.inner.frame()

crates/vm/src/coroutine.rs

@@ -32,7 +32,7 @@ pub struct Coro {
     // code
     // _weakreflist
     name: PyMutex<PyStrRef>,
-    // qualname
+    qualname: PyMutex<PyStrRef>,
     exception: PyMutex<Option<PyBaseExceptionRef>>, // exc_state
 }
 
@@ -48,13 +48,14 @@ fn gen_name(jen: &PyObject, vm: &VirtualMachine) -> &'static str {
 }
 
 impl Coro {
-    pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
+    pub fn new(frame: FrameRef, name: PyStrRef, qualname: PyStrRef) -> Self {
         Self {
             frame,
             closed: AtomicCell::new(false),
             running: AtomicCell::new(false),
             exception: PyMutex::default(),
             name: PyMutex::new(name),
+            qualname: PyMutex::new(qualname),
         }
     }
 
@@ -188,6 +189,14 @@ impl Coro {
         *self.name.lock() = name;
     }
 
+    pub fn qualname(&self) -> PyStrRef {
+        self.qualname.lock().clone()
+    }
+
+    pub fn set_qualname(&self, qualname: PyStrRef) {
+        *self.qualname.lock() = qualname;
+    }
+
     pub fn repr(&self, jen: &PyObject, id: usize, vm: &VirtualMachine) -> String {
         format!(
             "<{} object {} at {:#x}>",

crates/vm/src/protocol/object.rs

@@ -4,8 +4,8 @@
 use crate::{
     AsObject, Py, PyObject, PyObjectRef, PyRef, PyResult, TryFromObject, VirtualMachine,
     builtins::{
-        PyAsyncGen, PyBytes, PyDict, PyDictRef, PyGenericAlias, PyInt, PyList, PyStr, PyTuple,
-        PyTupleRef, PyType, PyTypeRef, PyUtf8Str, pystr::AsPyStr,
+        PyBytes, PyDict, PyDictRef, PyGenericAlias, PyInt, PyList, PyStr, PyTuple, PyTupleRef,
+        PyType, PyTypeRef, PyUtf8Str, pystr::AsPyStr,
     },
     common::{hash::PyHash, str::to_ascii},
     convert::{ToPyObject, ToPyResult},
@@ -87,11 +87,37 @@ impl PyObject {
 
     // PyObject *PyObject_GetAIter(PyObject *o)
     pub fn get_aiter(&self, vm: &VirtualMachine) -> PyResult {
-        if self.downcastable::<PyAsyncGen>() {
-            vm.call_special_method(self, identifier!(vm, __aiter__), ())
-        } else {
-            Err(vm.new_type_error("wrong argument type"))
+        use crate::builtins::PyCoroutine;
+
+        // Check if object has __aiter__ method
+        let aiter_method = self.class().get_attr(identifier!(vm, __aiter__));
+        let Some(_aiter_method) = aiter_method else {
+            return Err(vm.new_type_error(format!(
+                "'{}' object is not an async iterable",
+                self.class().name()
+            )));
+        };
+
+        // Call __aiter__
+        let iterator = vm.call_special_method(self, identifier!(vm, __aiter__), ())?;
+
+        // Check that __aiter__ did not return a coroutine
+        if iterator.downcast_ref::<PyCoroutine>().is_some() {
+            return Err(vm.new_type_error(
+                "'async_iterator' object cannot be interpreted as an async iterable; \
+                perhaps you forgot to call aiter()?",
+            ));
         }
+
+        // Check that the result is an async iterator (has __anext__)
+        if !iterator.class().has_attr(identifier!(vm, __anext__)) {
+            return Err(vm.new_type_error(format!(
+                "'{}' object is not an async iterator",
+                iterator.class().name()
+            )));
+        }
+
+        Ok(iterator)
     }
 
     pub fn has_attr<'a>(&self, attr_name: impl AsPyStr<'a>, vm: &VirtualMachine) -> PyResult<bool> {

crates/vm/src/stdlib/builtins.rs

@@ -540,12 +540,23 @@ mod builtins {
         default_value: OptionalArg<PyObjectRef>,
         vm: &VirtualMachine,
     ) -> PyResult {
+        use crate::builtins::asyncgenerator::PyAnextAwaitable;
+
+        // Check if object is an async iterator (has __anext__ method)
+        if !aiter.class().has_attr(identifier!(vm, __anext__)) {
+            return Err(vm.new_type_error(format!(
+                "'{}' object is not an async iterator",
+                aiter.class().name()
+            )));
+        }
+
         let awaitable = vm.call_method(&aiter, "__anext__", ())?;
 
-        if default_value.is_missing() {
-            Ok(awaitable)
+        if let OptionalArg::Present(default) = default_value {
+            Ok(PyAnextAwaitable::new(awaitable, default)
+                .into_ref(&vm.ctx)
+                .into())
         } else {
-            // TODO: Implement CPython like PyAnextAwaitable to properly handle the default value.
             Ok(awaitable)
         }
     }

crates/vm/src/types/zoo.rs

@@ -21,6 +21,7 @@ pub struct TypeZoo {
     pub async_generator_asend: &'static Py<PyType>,
     pub async_generator_athrow: &'static Py<PyType>,
     pub async_generator_wrapped_value: &'static Py<PyType>,
+    pub anext_awaitable: &'static Py<PyType>,
     pub bytes_type: &'static Py<PyType>,
     pub bytes_iterator_type: &'static Py<PyType>,
     pub bytearray_type: &'static Py<PyType>,
@@ -139,6 +140,7 @@ impl TypeZoo {
             async_generator_athrow: asyncgenerator::PyAsyncGenAThrow::init_builtin_type(),
             async_generator_wrapped_value:
                 asyncgenerator::PyAsyncGenWrappedValue::init_builtin_type(),
+            anext_awaitable: asyncgenerator::PyAnextAwaitable::init_builtin_type(),
             bound_method_type: function::PyBoundMethod::init_builtin_type(),
             builtin_function_or_method_type: builtin_func::PyNativeFunction::init_builtin_type(),
             builtin_method_type: builtin_func::PyNativeMethod::init_builtin_type(),