crates/vm/src/frame.rs (1)

1008-1032: Allow generator subclasses with __await__ to be awaited.

The current code uses downcast_ref::<PyGenerator>() which matches any generator, including subclasses, then immediately returns a TypeError if the ITERABLE_COROUTINE flag is absent. This blocks generator subclasses that implement __await__, preventing the fallback __await__ lookup.

Only generators that fail to have the ITERABLE_COROUTINE flag AND lack an __await__ method should error. The check should be restructured to only apply the generator-specific logic when the flag is actually set, allowing non-flagged generators to fall through to the __await__ fallback.

🔧 Suggested fix (preserve __await__ fallback for non-flagged generators)

-                } else if let Some(generator) = awaited_obj.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(bytecode::CodeFlags::ITERABLE_COROUTINE)
-                    {
-                        awaited_obj
-                    } else {
-                        return Err(vm.new_type_error(format!(
-                            "object {} can't be used in 'await' expression",
-                            awaited_obj.class().name(),
-                        )));
-                    }
-                } else {
+                } else if awaited_obj
+                    .downcast_ref::<PyGenerator>()
+                    .is_some_and(|generator| {
+                        generator
+                            .as_coro()
+                            .frame()
+                            .code
+                            .flags
+                            .contains(bytecode::CodeFlags::ITERABLE_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()
                         )));
                     }
                     result
                 };