From 4a29231b3a6c97bfdac0302c3a746b35f8918e5e Mon Sep 17 00:00:00 2001
From: snowapril <sinjihng@gmail.com>
Date: Fri, 1 Oct 2021 12:24:44 +0900
Subject: [PATCH] implement `PyIter_Check` on `PyIter`

Signed-off-by: snowapril <sinjihng@gmail.com>
---
 vm/src/protocol/iter.rs | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs
index e19f7e525..999e7baa5 100644
--- a/vm/src/protocol/iter.rs
+++ b/vm/src/protocol/iter.rs
@@ -18,6 +18,11 @@ impl PyIter<PyObjectRef> {
     pub fn into_object(self) -> PyObjectRef {
         self.0
     }
+    pub fn check(obj: &PyObjectRef) -> bool {
+        obj.class()
+            .mro_find_map(|x| x.slots.iternext.load())
+            .is_some()
+    }
 }
 
 impl<T> PyIter<T>
@@ -83,18 +88,16 @@ impl TryFromObject for PyIter<PyObjectRef> {
         };
         if let Some(getiter) = getiter {
             let iter = getiter(iter_target, vm)?;
-            let cls = iter.class();
-            let is_iter = cls.iter_mro().any(|x| x.slots.iternext.load().is_some());
-            if is_iter {
-                drop(cls);
+            if PyIter::check(&iter) {
                 Ok(Self(iter))
             } else {
                 Err(vm.new_type_error(format!(
                     "iter() returned non-iterator of type '{}'",
-                    cls.name()
+                    iter.class().name()
                 )))
             }
         } else {
+            // TODO: __getitem__ method lookup must be replaced by sequence protocol checking
             vm.get_method_or_type_error(iter_target.clone(), "__getitem__", || {
                 format!("'{}' object is not iterable", iter_target.class().name())
             })?;