Handle CO_ITERABLE_COROUTINE flag in anext_awaitable

moreal · claude · moreal · commit ae3a2eb8961e · 2026-01-22T17:52:26.000+09:00
Generators decorated with @types.coroutine have the CO_ITERABLE_COROUTINE
flag and can be used in await expressions. The PyAnextAwaitable's
get_awaitable_iter method was missing this check, causing builtin anext()
to fail with "object generator can't be used in 'await' expression" when
used with such generators.

Co-Authored-By: Claude Opus 4.5 &lt;noreply@anthropic.com&gt;
diff --git a/crates/vm/src/builtins/asyncgenerator.rs b/crates/vm/src/builtins/asyncgenerator.rs
@@ -1,4 +1,4 @@
-use super::{PyCode, PyGenericAlias, PyStrRef, PyType, PyTypeRef};
+use super::{PyCode, PyGenerator, PyGenericAlias, PyStrRef, PyType, PyTypeRef};
 use crate::{
     AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
     builtins::PyBaseExceptionRef,
@@ -592,6 +592,24 @@ impl PyAnextAwaitable {
         let awaitable = if wrapped.class().is(vm.ctx.types.coroutine_type) {
             // Coroutine - get __await__ later
             wrapped.clone()
+        } else if let Some(generator) = wrapped.downcast_ref::<PyGenerator>() {
+            // Generator with CO_ITERABLE_COROUTINE flag can be awaited
+            // (e.g., generators decorated with @types.coroutine)
+            if generator
+                .as_coro()
+                .frame()
+                .code
+                .flags
+                .contains(crate::bytecode::CodeFlags::ITERABLE_COROUTINE)
+            {
+                // Return the generator itself as the iterator
+                return Ok(wrapped.clone());
+            } else {
+                return Err(vm.new_type_error(format!(
+                    "object {} can't be used in 'await' expression",
+                    wrapped.class().name()
+                )));
+            }
         } else {
             // Try to get __await__ method
             if let Some(await_method) = vm.get_method(wrapped.clone(), identifier!(vm, __await__)) {
