Skip to content

Commit 35a06bc

Browse files
committed
Add bytes.__iter__
1 parent aaf0eab commit 35a06bc

2 files changed

Lines changed: 29 additions & 0 deletions

File tree

vm/src/obj/objbytes.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ pub fn init(context: &PyContext) {
2020
context.set_attr(bytes_type, "__new__", context.new_rustfunc(bytes_new));
2121
context.set_attr(bytes_type, "__repr__", context.new_rustfunc(bytes_repr));
2222
context.set_attr(bytes_type, "__len__", context.new_rustfunc(bytes_len));
23+
context.set_attr(bytes_type, "__iter__", context.new_rustfunc(bytes_iter))
2324
}
2425

2526
fn bytes_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -101,3 +102,17 @@ fn bytes_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
101102
let data = String::from_utf8(value.to_vec()).unwrap();
102103
Ok(vm.new_str(format!("b'{}'", data)))
103104
}
105+
106+
fn bytes_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
107+
arg_check!(vm, args, required = [(obj, Some(vm.ctx.bytes_type()))]);
108+
109+
let iter_obj = PyObject::new(
110+
PyObjectPayload::Iterator {
111+
position: 0,
112+
iterated_obj: obj.clone(),
113+
},
114+
vm.ctx.iter_type(),
115+
);
116+
117+
Ok(iter_obj)
118+
}

vm/src/obj/objiter.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,20 @@ fn iter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
131131
Err(stop_iteration)
132132
}
133133
}
134+
135+
PyObjectPayload::Bytes { ref value } => {
136+
if *position < value.len() {
137+
let obj_ref = vm.ctx.new_int(value[*position].to_bigint().unwrap());
138+
*position += 1;
139+
Ok(obj_ref)
140+
} else {
141+
let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
142+
let stop_iteration =
143+
vm.new_exception(stop_iteration_type, "End of iterator".to_string());
144+
Err(stop_iteration)
145+
}
146+
}
147+
134148
_ => {
135149
panic!("NOT IMPL");
136150
}

0 commit comments

Comments
 (0)