From 35a06bc42803b53da607d73ed8b0b0cf976fe233 Mon Sep 17 00:00:00 2001
From: Homer McMillan <him16@my.fsu.edu>
Date: Tue, 5 Feb 2019 21:44:14 -0500
Subject: [PATCH] Add bytes.__iter__

---
 vm/src/obj/objbytes.rs | 15 +++++++++++++++
 vm/src/obj/objiter.rs  | 14 ++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
index f700b8bdb..8b0f60bd5 100644
--- a/vm/src/obj/objbytes.rs
+++ b/vm/src/obj/objbytes.rs
@@ -20,6 +20,7 @@ pub fn init(context: &PyContext) {
     context.set_attr(bytes_type, "__new__", context.new_rustfunc(bytes_new));
     context.set_attr(bytes_type, "__repr__", context.new_rustfunc(bytes_repr));
     context.set_attr(bytes_type, "__len__", context.new_rustfunc(bytes_len));
+    context.set_attr(bytes_type, "__iter__", context.new_rustfunc(bytes_iter))
 }
 
 fn bytes_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -101,3 +102,17 @@ fn bytes_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let data = String::from_utf8(value.to_vec()).unwrap();
     Ok(vm.new_str(format!("b'{}'", data)))
 }
+
+fn bytes_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, Some(vm.ctx.bytes_type()))]);
+
+    let iter_obj = PyObject::new(
+        PyObjectPayload::Iterator {
+            position: 0,
+            iterated_obj: obj.clone(),
+        },
+        vm.ctx.iter_type(),
+    );
+
+    Ok(iter_obj)
+}
diff --git a/vm/src/obj/objiter.rs b/vm/src/obj/objiter.rs
index 53ed76804..86eb79c77 100644
--- a/vm/src/obj/objiter.rs
+++ b/vm/src/obj/objiter.rs
@@ -131,6 +131,20 @@ fn iter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
                     Err(stop_iteration)
                 }
             }
+
+            PyObjectPayload::Bytes { ref value } => {
+                if *position < value.len() {
+                    let obj_ref = vm.ctx.new_int(value[*position].to_bigint().unwrap());
+                    *position += 1;
+                    Ok(obj_ref)
+                } else {
+                    let stop_iteration_type = vm.ctx.exceptions.stop_iteration.clone();
+                    let stop_iteration =
+                        vm.new_exception(stop_iteration_type, "End of iterator".to_string());
+                    Err(stop_iteration)
+                }
+            }
+
             _ => {
                 panic!("NOT IMPL");
             }