8000 iter with slot-wrapper · RustPython/RustPython@47fcdd9 · GitHub
[go: up one dir, main page]
Skip to content

Commit 47fcdd9

Browse files
youknowoneyouknowone
committed
iter with slot-wrapper
1 parent c2a7393 commit 47fcdd9

File tree

3 files changed

+69
-59
lines changed

3 files changed

+69
-59
lines changed

‎crates/vm/src/builtins/descriptor.rs‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,11 @@ use crate::{
44
builtins::{PyTypeRef, builtin_func::PyNativeMethod, type_},
55
class::PyClassImpl,
66
common::hash::PyHash,
7+
convert::ToPyResult,
78
function::{FuncArgs, PyMethodDef, PyMethodFlags, PySetterValue},
89
types::{
9-
Callable, Comparable, GetDescriptor, HashFunc, Hashable, InitFunc, PyComparisonOp,
10-
Representable, StringifyFunc,
10+
Callable, Comparable, GetDescriptor, HashFunc, Hashable, InitFunc, IterFunc, IterNextFunc,
11+
PyComparisonOp, Representable, StringifyFunc,
1112
},
1213
};
1314
use rustpython_common::lock::PyRwLock;
@@ -399,6 +400,8 @@ pub enum SlotFunc {
399400
Init(InitFunc),
400401
Hash(HashFunc),
401402
Repr(StringifyFunc),
403+
Iter(IterFunc),
404+
IterNext(IterNextFunc),
402405
}
403406

404407
impl std::fmt::Debug for SlotFunc {
@@ -407,6 +410,8 @@ impl std::fmt::Debug for SlotFunc {
407410
SlotFunc::Init(_) => write!(f, "SlotFunc::Init(...)"),
408411
SlotFunc::Hash(_) => write!(f, "SlotFunc::Hash(...)"),
409412
SlotFunc::Repr(_) => write!(f, "SlotFunc::Repr(...)"),
413+
SlotFunc::Iter(_) => write!(f, "SlotFunc::Iter(...)"),
414+
SlotFunc::IterNext(_) => write!(f, "SlotFunc::IterNext(...)"),
410415
}
411416
}
412417
}
@@ -437,6 +442,22 @@ impl SlotFunc {
437442
let s = func(&obj, vm)?;
438443
Ok(s.into())
439444
}
445+
SlotFunc::Iter(func) => {
446+
if !args.args.is_empty() || !args.kwargs.is_empty() {
447+
return Err(
448+
vm.new_type_error("__iter__() takes no arguments (1 given)".to_owned())
449+
);
450+
}
451+
func(obj, vm)
452+
}
453+
SlotFunc::IterNext(func) => {
454+
if !args.args.is_empty() || !args.kwargs.is_empty() {
455+
return Err(
456+
vm.new_type_error("__next__() takes no arguments (1 given)".to_owned())
457+
);
458+
}
459+
func(&obj, vm).to_pyresult(vm)
460+
}
440461
}
441462
}
442463
}

‎crates/vm/src/class.rs‎

Lines changed: 42 additions & 44 deletions
9E81
Original file line numberDiff line numberDiff line change
@@ -139,52 +139,50 @@ pub trait PyClassImpl: PyClassDef {
139139
}
140140
}
141141

142-
// Add __init__ slot wrapper if slot exists and not already in dict
143-
if let Some(init_func) = class.slots.init.load() {
144-
let init_name = identifier!(ctx, __init__);
145-
if !class.attributes.read().contains_key(init_name) {
146-
let wrapper = PySlotWrapper {
147-
typ: class,
148-
name: ctx.intern_str("__init__"),
149-
wrapped: SlotFunc::Init(init_func),
150-
doc: Some("Initialize self. See help(type(self)) for accurate signature."),
151-
};
152-
class.set_attr(init_name, wrapper.into_ref(ctx).into());
153-
}
154-
}
155-
156-
// Add __hash__ slot wrapper if slot exists and not already in dict
157-
// Note: hash_not_implemented is handled separately (sets __hash__ = None)
158-
if let Some(hash_func) = class.slots.hash.load()
159-
&& hash_func as usize != hash_not_implemented as usize
160-
{
161-
let hash_name = identifier!(ctx, __hash__);
162-
if !class.attributes.read().contains_key(hash_name) {
163-
let wrapper = PySlotWrapper {
164-
typ: class,
165-
name: ctx.intern_str("__hash__"),
166-
wrapped: SlotFunc::Hash(hash_func),
167-
doc: Some("Return hash(self)."),
168-
};
169-
class.set_attr(hash_name, wrapper.into_ref(ctx).into());
170-
}
171-
}
172-
173-
if class.slots.hash.load().map_or(0, |h| h as usize) == hash_not_implemented as usize {
174-
class.set_attr(ctx.names.__hash__, ctx.none.clone().into());
142+
// Add slot wrappers for slots that exist and are not already in dict
143+
// This mirrors CPython's add_operators() in typeobject.c
144+
macro_rules! add_slot_wrapper {
145+
($slot:ident, $name:ident, $variant:ident, $doc:expr) => {
146+
if let Some(func) = class.slots.$slot.load() {
147+
let attr_name = identifier!(ctx, $name);
148+
if !class.attributes.read().contains_key(attr_name) {
149+
let wrapper = PySlotWrapper {
150+
typ: class,
151+
name: ctx.intern_str(stringify!($name)),
152+
wrapped: SlotFunc::$variant(func),
153+
doc: Some($doc),
154+
};
155+
class.set_attr(attr_name, wrapper.into_ref( 6D3F ctx).into());
156+
}
157+
}
158+
};
175159
}
176160

177-
// Add __repr__ slot wrapper if slot exists and not already in dict
178-
if let Some(repr_func) = class.slots.repr.load() {
179-
let repr_name = identifier!(ctx, __repr__);
180-
if !class.attributes.read().contains_key(repr_name) {
181-
let wrapper = PySlotWrapper {
182-
typ: class,
183-
name: ctx.intern_str("__repr__"),
184-
wrapped: SlotFunc::Repr(repr_func),
185-
doc: Some("Return repr(self)."),
186-
};
187-
class.set_attr(repr_name, wrapper.into_ref(ctx).into());
161+
add_slot_wrapper!(
162+
init,
163+
__init__,
164+
Init,
165+
"Initialize self. See help(type(self)) for accurate signature."
166+
);
167+
add_slot_wrapper!(repr, __repr__, Repr, "Return repr(self).");
168+
add_slot_wrapper!(iter, __iter__, Iter, "Implement iter(self).");
169+
add_slot_wrapper!(iternext, __next__, IterNext, "Implement next(self).");
170+
171+
// __hash__ needs special handling: hash_not_implemented sets __hash__ = None
172+
if let Some(hash_func) = class.slots.hash.load() {
173+
if hash_func as usize == hash_not_implemented as usize {
174+
class.set_attr(ctx.names.__hash__, ctx.none.clone().into());
175+
} else {
176+
let hash_name = identifier!(ctx, __hash__);
177+
if !class.attributes.read().contains_key(hash_name) {
178+
let wrapper = PySlotWrapper {
179+
typ: class,
180+
name: ctx.intern_str("__hash__"),
181+
wrapped: SlotFunc::Hash(hash_func),
182+
doc: Some("Return hash(self)."),
183+
};
184+
class.set_attr(hash_name, wrapper.into_ref(ctx).into());
185+
}
188186
}
189187
}
190188

‎crates/vm/src/types/slot.rs‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use crate::{
66
builtins::{PyInt, PyStr, PyStrInterned, PyStrRef, PyType, PyTypeRef, type_::PointerSlot},
77
bytecode::ComparisonOperator,
88
common::hash::PyHash,
9-
convert::{ToPyObject, ToPyResult},
9+
convert::ToPyObject,
1010
function::{
1111
Either, FromArgs, FuncArgs, OptionalArg, PyComparisonValue, PyMethodDef, PySetterValue,
1212
},
@@ -1435,10 +1435,7 @@ pub trait Iterable: PyPayload {
14351435
Self::iter(zelf, vm)
14361436
}
14371437

1438-
#[pymethod]
1439-
fn __iter__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
1440-
Self::slot_iter(zelf, vm)
1441-
}
1438+
// __iter__ is exposed via SlotFunc::Iter wrapper in extend_class()
14421439

14431440
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult;
14441441

@@ -1458,11 +1455,7 @@ pub trait IterNext: PyPayload + Iterable {
14581455

14591456
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn>;
14601457

1461-
#[inline]
1462-
#[pymethod]
1463-
fn __next__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
1464-
Self::slot_iternext(&zelf, vm).to_pyresult(vm)
1465-
}
1458+
// __next__ is exposed via SlotFunc::IterNext wrapper in extend_class()
14661459
}
14671460

14681461
pub trait SelfIter: PyPayload {}
@@ -1477,9 +1470,7 @@ where
14771470
unreachable!("slot must be overridden for {}", repr.as_str());
14781471
}
14791472

1480-
fn __iter__(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult {
1481-
self_iter(zelf, vm)
1482-
}
1473+
// __iter__ is exposed via SlotFunc::Iter wrapper in extend_class()
14831474

14841475
#[cold]
14851476
fn iter(_zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult {

0 commit comments

Comments
 (0)
0