From e1f689c5651749af68eeb214d3d9e1085034824c Mon Sep 17 00:00:00 2001
From: Kangzhi Shi <shikangzhi@gmail.com>
Date: Fri, 1 Oct 2021 21:16:43 +0200
Subject: [PATCH] introduce PyIterReturn::from_getitem_result

---
 vm/src/builtins/enumerate.rs | 13 +------------
 vm/src/builtins/iter.rs      | 15 ++-------------
 vm/src/protocol/iter.rs      | 13 +++++++++++++
 3 files changed, 16 insertions(+), 25 deletions(-)

diff --git a/vm/src/builtins/enumerate.rs b/vm/src/builtins/enumerate.rs
index 4922e491d..115631753 100644
--- a/vm/src/builtins/enumerate.rs
+++ b/vm/src/builtins/enumerate.rs
@@ -1,6 +1,5 @@
 use super::{IterStatus, PositionIterInternal, PyIntRef, PyTypeRef};
 use crate::common::lock::{PyMutex, PyRwLock};
-use crate::TypeProtocol;
 use crate::{
     function::OptionalArg,
     protocol::{PyIter, PyIterReturn},
@@ -115,17 +114,7 @@ impl SlotIterator for PyReverseSequenceIterator {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         zelf.internal
             .lock()
-            .rev_next(|obj, pos| match obj.get_item(pos, vm) {
-                Ok(ret) => Ok(PyIterReturn::Return(ret)),
-                Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => {
-                    Ok(PyIterReturn::StopIteration(None))
-                }
-                Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => {
-                    let args = e.get_arg(0);
-                    Ok(PyIterReturn::StopIteration(args))
-                }
-                Err(e) => Err(e),
-            })
+            .rev_next(|obj, pos| PyIterReturn::from_getitem_result(obj.get_item(pos, vm), vm))
     }
 }
 
diff --git a/vm/src/builtins/iter.rs b/vm/src/builtins/iter.rs
index de39b23e3..6e9d64e49 100644
--- a/vm/src/builtins/iter.rs
+++ b/vm/src/builtins/iter.rs
@@ -7,8 +7,7 @@ use crate::{
     function::ArgCallable,
     protocol::PyIterReturn,
     slots::{IteratorIterable, SlotIterator},
-    ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
-    VirtualMachine,
+    ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
 };
 use rustpython_common::{
     lock::{PyMutex, PyRwLock, PyRwLockUpgradableReadGuard},
@@ -211,17 +210,7 @@ impl SlotIterator for PySequenceIterator {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
         zelf.internal
             .lock()
-            .next(|obj, pos| match obj.get_item(pos, vm) {
-                Ok(ret) => Ok(PyIterReturn::Return(ret)),
-                Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => {
-                    Ok(PyIterReturn::StopIteration(None))
-                }
-                Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => {
-                    let args = e.get_arg(0);
-                    Ok(PyIterReturn::StopIteration(args))
-                }
-                Err(e) => Err(e),
-            })
+            .next(|obj, pos| PyIterReturn::from_getitem_result(obj.get_item(pos, vm), vm))
     }
 }
 
diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs
index 967333bff..9d081249f 100644
--- a/vm/src/protocol/iter.rs
+++ b/vm/src/protocol/iter.rs
@@ -136,6 +136,19 @@ impl PyIterReturn {
             Err(err) => Err(err),
         }
     }
+    pub fn from_getitem_result(result: PyResult, vm: &VirtualMachine) -> PyResult<Self> {
+        match result {
+            Ok(obj) => Ok(Self::Return(obj)),
+            Err(err) if err.isinstance(&vm.ctx.exceptions.index_error) => {
+                Ok(Self::StopIteration(None))
+            }
+            Err(err) if err.isinstance(&vm.ctx.exceptions.stop_iteration) => {
+                let args = err.get_arg(0);
+                Ok(Self::StopIteration(args))
+            }
+            Err(err) => Err(err),
+        }
+    }
 }
 
 impl IntoPyResult for PyIterReturn {