diff --git a/vm/src/protocol/object.rs b/vm/src/protocol/object.rs index 639e24dda..0b87bbf9d 100644 --- a/vm/src/protocol/object.rs +++ b/vm/src/protocol/object.rs @@ -3,8 +3,8 @@ use crate::{ builtins::{ - pystr::AsPyStr, PyBytes, PyDict, PyDictRef, PyGenericAlias, PyInt, PyList, PyStr, PyStrRef, - PyTuple, PyTupleRef, PyType, PyTypeRef, + pystr::AsPyStr, PyAsyncGen, PyBytes, PyDict, PyDictRef, PyGenericAlias, PyInt, PyList, + PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef, }, bytesinner::ByteInnerNewOptions, common::{hash::PyHash, str::to_ascii}, @@ -92,6 +92,13 @@ impl PyObject { } // PyObject *PyObject_GetAIter(PyObject *o) + pub fn get_aiter(&self, vm: &VirtualMachine) -> PyResult { + if self.payload_is::() { + vm.call_special_method(self, identifier!(vm, __aiter__), ()) + } else { + Err(vm.new_type_error("wrong argument type".to_owned())) + } + } pub fn has_attr<'a>(&self, attr_name: impl AsPyStr<'a>, vm: &VirtualMachine) -> PyResult { self.get_attr(attr_name, vm).map(|o| !vm.is_none(&o)) diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs index 835f4152e..4cc3a5363 100644 --- a/vm/src/stdlib/builtins.rs +++ b/vm/src/stdlib/builtins.rs @@ -9,7 +9,6 @@ pub use builtins::{ascii, print, reversed}; mod builtins { use crate::{ builtins::{ - asyncgenerator::PyAsyncGen, enumerate::PyReverseSequenceIterator, function::{PyCellRef, PyFunction}, int::PyIntRef, @@ -459,11 +458,7 @@ mod builtins { #[pyfunction] fn aiter(iter_target: PyObjectRef, vm: &VirtualMachine) -> PyResult { - if iter_target.payload_is::() { - vm.call_special_method(&iter_target, identifier!(vm, __aiter__), ()) - } else { - Err(vm.new_type_error("wrong argument type".to_owned())) - } + iter_target.get_aiter(vm) } #[pyfunction]