Add object protocol correspoinding to PyObject_GetAIter (#5090)

dannasman · web-flow · commit 4e6172b99d64 · 2023-10-16T10:39:29.000-07:00
diff --git 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::<PyAsyncGen>() {
+            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<bool> {
         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
@@ -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::<PyAsyncGen>() {
-            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]
