GetANext

youknowone · youknowone · commit acd7c225e162 · 2024-04-24T19:03:55.000+09:00
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -995,11 +995,33 @@ impl ExecutingFrame<'_> {
             }
             bytecode::Instruction::GetANext => {
                 let aiter = self.last_value();
-                let awaitable = vm.call_special_method(&aiter, identifier!(vm, __anext__), ())?;
-                let awaitable = if awaitable.payload_is::<PyCoroutine>() {
-                    awaitable
+                let awaitable = if aiter.class().is(vm.ctx.types.async_generator) {
+                    vm.call_special_method(&aiter, identifier!(vm, __anext__), ())?
                 } else {
-                    vm.call_special_method(&awaitable, identifier!(vm, __await__), ())?
+                    assert!(false);
+                    // if !hasattr(aiter, __anext__): raise TypeError("'async for' requires in iterator...")
+                    let next_iter =
+                        vm.call_special_method(&aiter, identifier!(vm, __anext__), ())?;
+
+                    // TODO: _PyCoro_GetAwaitableIter
+                    fn get_awaitable_iter(next_iter: &PyObject, vm: &VirtualMachine) -> PyResult {
+                        let gen_is_coroutine = |_| {
+                            // TODO:
+                            true
+                        };
+                        if next_iter.class().is(vm.ctx.types.coroutine_type)
+                            || gen_is_coroutine(next_iter)
+                        {
+                            return Ok(next_iter.to_owned());
+                        }
+                        vm.call_special_method(next_iter, identifier!(vm, __await__), ())
+                    }
+                    get_awaitable_iter(&next_iter, vm).map_err(|_| {
+                        vm.new_type_error(format!(
+                            "'async for' received an invalid object from __anext__: {:.200}",
+                            next_iter.class().name()
+                        ))
+                    })?
                 };
                 self.push_value(awaitable);
                 Ok(None)
