enumerate

OddCoincidence · OddCoincidence · commit a2951f8a8b19 · 2019-03-08T18:29:58.000-08:00
diff --git a/vm/src/obj/objenumerate.rs b/vm/src/obj/objenumerate.rs
@@ -3,11 +3,26 @@ use std::ops::AddAssign;
 
 use super::objint;
 use super::objiter;
-use crate::pyobject::{PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyResult, TypeProtocol};
+use crate::pyobject::{
+    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectPayload2, PyObjectRef, PyResult,
+    TypeProtocol,
+};
 use crate::vm::VirtualMachine;
 use num_bigint::BigInt;
 use num_traits::Zero;
 
+#[derive(Debug)]
+pub struct PyEnumerate {
+    counter: RefCell<BigInt>,
+    iterator: PyObjectRef,
+}
+
+impl PyObjectPayload2 for PyEnumerate {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.enumerate_type()
+    }
+}
+
 fn enumerate_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -22,9 +37,11 @@ fn enumerate_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     };
     let iterator = objiter::get_iter(vm, iterable)?;
     Ok(PyObject::new(
-        PyObjectPayload::EnumerateIterator {
-            counter: RefCell::new(counter),
-            iterator,
+        PyObjectPayload::AnyRustValue {
+            value: Box::new(PyEnumerate {
+                counter: RefCell::new(counter),
+                iterator,
+            }),
         },
         cls.clone(),
     ))
@@ -37,10 +54,10 @@ fn enumerate_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(enumerate, Some(vm.ctx.enumerate_type()))]
     );
 
-    if let PyObjectPayload::EnumerateIterator {
+    if let Some(PyEnumerate {
         ref counter,
         ref iterator,
-    } = enumerate.payload
+    }) = enumerate.payload()
     {
         let next_obj = objiter::call_next(vm, iterator)?;
         let result = vm
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -1504,10 +1504,6 @@ pub enum PyObjectPayload {
         position: Cell<usize>,
         iterated_obj: PyObjectRef,
     },
-    EnumerateIterator {
-        counter: RefCell<BigInt>,
-        iterator: PyObjectRef,
-    },
     Slice {
         start: Option<BigInt>,
         stop: Option<BigInt>,
@@ -1560,7 +1556,6 @@ impl fmt::Debug for PyObjectPayload {
             PyObjectPayload::MemoryView { ref obj } => write!(f, "bytes/bytearray {:?}", obj),
             PyObjectPayload::WeakRef { .. } => write!(f, "weakref"),
             PyObjectPayload::Iterator { .. } => write!(f, "iterator"),
-            PyObjectPayload::EnumerateIterator { .. } => write!(f, "enumerate"),
             PyObjectPayload::Slice { .. } => write!(f, "slice"),
             PyObjectPayload::Function { .. } => write!(f, "function"),
             PyObjectPayload::Generator { .. } => write!(f, "generator"),
