Add iter(callable, sentinel)

coolreader18 · coolreader18 · commit fd3fcc244907 · 2020-04-08T10:10:56.000-05:00
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -29,8 +29,8 @@ use crate::obj::objstr::{PyString, PyStringRef};
 use crate::obj::objtype::{self, PyClassRef};
 use crate::pyhash;
 use crate::pyobject::{
-    Either, IdProtocol, ItemProtocol, PyIterable, PyObjectRef, PyResult, PyValue, TryFromObject,
-    TypeProtocol,
+    Either, IdProtocol, ItemProtocol, PyCallable, PyIterable, PyObjectRef, PyResult, PyValue,
+    TryFromObject, TypeProtocol,
 };
 use crate::readline::{Readline, ReadlineResult};
 use crate::scope::Scope;
@@ -404,8 +404,19 @@ fn builtin_issubclass(
     )
 }
 
-fn builtin_iter(iter_target: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-    objiter::get_iter(vm, &iter_target)
+fn builtin_iter(
+    iter_target: PyObjectRef,
+    sentinel: OptionalArg<PyObjectRef>,
+    vm: &VirtualMachine,
+) -> PyResult {
+    if let OptionalArg::Present(sentinel) = sentinel {
+        let callable = PyCallable::try_from_object(vm, iter_target)?;
+        Ok(objiter::PyCallableIterator::new(callable, sentinel)
+            .into_ref(vm)
+            .into_object())
+    } else {
+        objiter::get_iter(vm, &iter_target)
+    }
 }
 
 fn builtin_len(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
diff --git a/vm/src/obj/objiter.rs b/vm/src/obj/objiter.rs
@@ -10,7 +10,8 @@ use super::objsequence;
 use super::objtype::{self, PyClassRef};
 use crate::exceptions::PyBaseExceptionRef;
 use crate::pyobject::{
-    PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol,
+    PyCallable, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
+    TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 
@@ -201,6 +202,53 @@ pub fn seq_iter_method(obj: PyObjectRef) -> PySequenceIterator {
     }
 }
 
+#[pyclass(name = "callable_iterator")]
+#[derive(Debug)]
+pub struct PyCallableIterator {
+    callable: PyCallable,
+    sentinel: PyObjectRef,
+    done: Cell<bool>,
+}
+
+impl PyValue for PyCallableIterator {
+    fn class(vm: &VirtualMachine) -> PyClassRef {
+        vm.ctx.types.callable_iterator.clone()
+    }
+}
+
+#[pyimpl]
+impl PyCallableIterator {
+    pub fn new(callable: PyCallable, sentinel: PyObjectRef) -> Self {
+        Self {
+            callable,
+            sentinel,
+            done: Cell::new(false),
+        }
+    }
+
+    #[pymethod(magic)]
+    fn next(&self, vm: &VirtualMachine) -> PyResult {
+        if self.done.get() {
+            return Err(new_stop_iteration(vm));
+        }
+
+        let ret = self.callable.invoke(vec![], vm)?;
+
+        if vm.bool_eq(ret.clone(), self.sentinel.clone())? {
+            self.done.set(true);
+            Err(new_stop_iteration(vm))
+        } else {
+            Ok(ret)
+        }
+    }
+
+    #[pymethod(magic)]
+    fn iter(zelf: PyRef<Self>) -> PyRef<Self> {
+        zelf
+    }
+}
+
 pub fn init(context: &PyContext) {
     PySequenceIterator::extend_class(context, &context.types.iter_type);
+    PyCallableIterator::extend_class(context, &context.types.callable_iterator);
 }
diff --git a/vm/src/types.rs b/vm/src/types.rs
@@ -57,6 +57,7 @@ pub struct TypeZoo {
     pub bytearray_type: PyClassRef,
     pub bytearrayiterator_type: PyClassRef,
     pub bool_type: PyClassRef,
+    pub callable_iterator: PyClassRef,
     pub classmethod_type: PyClassRef,
     pub code_type: PyClassRef,
     pub coroutine_type: PyClassRef,
@@ -184,6 +185,7 @@ impl TypeZoo {
         let slice_type = create_type("slice", &type_type, &object_type);
         let mappingproxy_type = create_type("mappingproxy", &type_type, &object_type);
         let traceback_type = create_type("traceback", &type_type, &object_type);
+        let callable_iterator = create_type("callable_iterator", &type_type, &object_type);
 
         Self {
             async_generator,
@@ -196,6 +198,7 @@ impl TypeZoo {
             bytearrayiterator_type,
             bytes_type,
             bytesiterator_type,
+            callable_iterator,
             code_type,
             coroutine_type,
             coroutine_wrapper_type,
