more asyncgen impl

Author: youknowone

Lib/test/test_asyncgen.py

@@ -467,16 +467,12 @@ async def test_throw():
         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 +488,6 @@ async def __anext__(self):
                     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:
@@ -549,8 +543,6 @@ async def gen():
         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 +563,7 @@ async def call_with_kwarg():
         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 +634,7 @@ async def do_test():
         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

crates/vm/src/builtins/asyncgenerator.rs

@@ -424,8 +424,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.
+    /// This is equivalent to CPython's 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

@@ -156,6 +156,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/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(),