From 21fffc523e71e4f4661debc6067551c5dcf3ce69 Mon Sep 17 00:00:00 2001 From: Bas Schoenmaeckers Date: Wed, 3 Jun 2026 16:17:29 +0200 Subject: [PATCH] Add iter support to c-api --- crates/capi/src/abstract_.rs | 10 +++ crates/capi/src/abstract_/iter.rs | 129 ++++++++++++++++++++++++++++++ 2 files changed, 139 insertions(+) create mode 100644 crates/capi/src/abstract_/iter.rs diff --git a/crates/capi/src/abstract_.rs b/crates/capi/src/abstract_.rs index a677c98c7b4..01a3964f728 100644 --- a/crates/capi/src/abstract_.rs +++ b/crates/capi/src/abstract_.rs @@ -1,6 +1,7 @@ use crate::{PyObject, pystate::with_vm}; use alloc::slice; use core::ffi::c_int; +pub use iter::*; pub use mapping::*; pub use number::*; use rustpython_vm::builtins::{PyDict, PyStr, PyTuple}; @@ -8,6 +9,7 @@ use rustpython_vm::function::{FuncArgs, KwArgs, PosArgs}; use rustpython_vm::{AsObject, Py, PyObjectRef, PyResult, VirtualMachine}; pub use sequence::*; +mod iter; mod mapping; mod number; mod sequence; @@ -170,3 +172,11 @@ pub unsafe extern "C" fn PyObject_IsInstance(inst: *mut PyObject, cls: *mut PyOb inst.is_instance(cls, vm) }) } + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyObject_Size(obj: *mut PyObject) -> isize { + with_vm(|vm| { + let obj = unsafe { &*obj }; + obj.length(vm) + }) +} diff --git a/crates/capi/src/abstract_/iter.rs b/crates/capi/src/abstract_/iter.rs new file mode 100644 index 00000000000..1ba5bd04d19 --- /dev/null +++ b/crates/capi/src/abstract_/iter.rs @@ -0,0 +1,129 @@ +use crate::{PyObject, pystate::with_vm}; +use core::ffi::c_int; +use rustpython_vm::PyObjectRef; +use rustpython_vm::builtins::PyGenerator; +use rustpython_vm::protocol::{PyIter, PyIterReturn}; + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyIter_Check(obj: *mut PyObject) -> c_int { + with_vm(|_vm| Ok(PyIter::check(unsafe { &*obj }))) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyObject_GetIter(obj: *mut PyObject) -> *mut PyObject { + with_vm(|vm| { + let obj = unsafe { &*obj }; + obj.get_iter(vm).map(PyObjectRef::from) + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyIter_NextItem(iter: *mut PyObject, item: *mut *mut PyObject) -> c_int { + with_vm(|vm| { + unsafe { + *item = core::ptr::null_mut(); + } + + let iter = PyIter::new(unsafe { &*iter }); + match iter.next(vm)? { + PyIterReturn::Return(next_item) => { + unsafe { + *item = next_item.into_raw().as_ptr(); + }; + Ok(true) + } + PyIterReturn::StopIteration(_) => Ok(false), + } + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyIter_Next(iter: *mut PyObject) -> *mut PyObject { + with_vm(|vm| { + let iter = PyIter::new(unsafe { &*iter }); + match iter.next(vm)? { + PyIterReturn::Return(next_item) => Ok(next_item.into_raw().as_ptr()), + PyIterReturn::StopIteration(_) => Ok(core::ptr::null_mut()), + } + }) +} + +#[unsafe(no_mangle)] +pub unsafe extern "C" fn PyIter_Send( + iter: *mut PyObject, + arg: *mut PyObject, + presult: *mut *mut PyObject, +) -> c_int { + with_vm(|vm| { + unsafe { + *presult = core::ptr::null_mut(); + } + + let iter_obj = unsafe { &*iter }; + let arg_obj = unsafe { &*arg }; + + let ret = if vm.is_none(arg_obj) { + PyIter::new(iter_obj).next(vm)? + } else { + iter_obj + .try_downcast_ref::(vm)? + .as_coro() + .send(iter_obj, arg_obj.to_owned(), vm)? + }; + + match ret { + PyIterReturn::Return(next_item) => { + unsafe { + *presult = next_item.into_raw().as_ptr(); + }; + Ok(true) + } + PyIterReturn::StopIteration(ret_val) => { + let ret_val = ret_val.unwrap_or_else(|| vm.ctx.none()); + unsafe { + *presult = ret_val.into_raw().as_ptr(); + }; + Ok(false) + } + } + }) +} + +#[cfg(false)] +mod tests { + use pyo3::prelude::*; + use pyo3::types::{PyAnyMethods, PyIterator, PyList, PySendResult}; + + #[test] + fn next_item() { + Python::attach(|py| { + let list = PyList::new(py, [1, 2, 3]).unwrap(); + let iter = list.try_iter().unwrap(); + let items: Vec = iter.map(|x| x.unwrap().extract::().unwrap()).collect(); + assert_eq!(items, vec![1, 2, 3]); + }) + } + + #[test] + fn send_generator() { + Python::attach(|py| { + let generator = py + .eval(c"(x for x in (1, 2))", None, None) + .unwrap() + .cast_into::() + .unwrap(); + + let first = generator.send(py.None().bind(py)).unwrap(); + assert!(matches!( + first, + PySendResult::Next(value) if value.extract::().unwrap() == 1 + )); + + let second = generator.send(py.None().bind(py)).unwrap(); + assert!(matches!( + second, + PySendResult::Next(value) if value.extract::().unwrap() == 2 + )); + }) + } +}