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 { 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 PyIter @@ -83,18 +88,16 @@ impl TryFromObject for PyIter { }; 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()) })?;