From d0615a4ea3c8101af7a865d51053e843246effcf Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Mon, 4 Oct 2021 00:34:52 +0900
Subject: [PATCH] PyIter::iter

---
 vm/src/function/argument.rs |  9 +++------
 vm/src/protocol/iter.rs     | 10 ++++++++++
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/vm/src/function/argument.rs b/vm/src/function/argument.rs
index 9c19fbbab..d9fc1ac15 100644
--- a/vm/src/function/argument.rs
+++ b/vm/src/function/argument.rs
@@ -51,14 +51,11 @@ impl<T> ArgIterable<T> {
     /// This operation may fail if an exception is raised while invoking the
     /// `__iter__` method of the iterable object.
     pub fn iter<'a>(&self, vm: &'a VirtualMachine) -> PyResult<PyIterIter<'a, T>> {
-        let iter_obj = match self.iterfn {
+        let iter = PyIter::new(match self.iterfn {
             Some(f) => f(self.iterable.clone(), vm)?,
             None => PySequenceIterator::new(self.iterable.clone()).into_object(vm),
-        };
-
-        let length_hint = vm.length_hint(iter_obj.clone())?;
-
-        Ok(PyIterIter::new(vm, PyIter::new(iter_obj), length_hint))
+        });
+        iter.iter(vm)
     }
 }
 
diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs
index d64943d08..e08755652 100644
--- a/vm/src/protocol/iter.rs
+++ b/vm/src/protocol/iter.rs
@@ -50,6 +50,16 @@ where
         };
         iternext(self.0.borrow(), vm)
     }
+
+    pub fn iter<'a, U>(&self, vm: &'a VirtualMachine) -> PyResult<PyIterIter<'a, U>> {
+        let obj = self.as_object();
+        let length_hint = vm.length_hint(obj.clone())?;
+        Ok(PyIterIter::new(
+            vm,
+            PyIter::<PyObjectRef>::new(obj.clone()),
+            length_hint,
+        ))
+    }
 }
 
 impl<T> Borrow<PyObjectRef> for PyIter<T>