From 89e8341d87e70fa3e05504b1895f046bcb8cf1b9 Mon Sep 17 00:00:00 2001
From: wdsmini <wdsmini@wdsmini.local>
Date: Fri, 20 Mar 2026 11:28:08 +0800
Subject: [PATCH] fix(bytes): make bytes object iterable

Add __iter__ method to bytes type and implement bytes_iterator type.
This allows iterating over bytes objects to get individual byte values
as integers (0-255), matching Python's behavior.

Changes:
- include/pocketpy/objects/iterator.h: Add bytes_iterator struct
- include/pocketpy/pocketpy.h: Add tp_bytes_iterator type constant
- include/pocketpy/interpreter/vm.h: Declare pk_bytes_iterator__register
- src/bindings/py_str.c: Implement bytes__iter__ and bytes_iterator__next__
- src/interpreter/vm.c: Register bytes_iterator type

Fixes #450
---
 include/pocketpy/interpreter/vm.h   |  1 +
 include/pocketpy/objects/iterator.h |  6 ++++++
 include/pocketpy/pocketpy.h         |  1 +
 src/bindings/py_str.c               | 30 +++++++++++++++++++++++++++++
 src/interpreter/vm.c                |  1 +
 5 files changed, 39 insertions(+)

diff --git a/include/pocketpy/interpreter/vm.h b/include/pocketpy/interpreter/vm.h
index 1b3126ac..769d95c6 100644
--- a/include/pocketpy/interpreter/vm.h
+++ b/include/pocketpy/interpreter/vm.h
@@ -133,6 +133,7 @@ void pk_number__register();
 py_Type pk_str__register();
 py_Type pk_str_iterator__register();
 py_Type pk_bytes__register();
+py_Type pk_bytes_iterator__register();
 py_Type pk_dict__register();
 py_Type pk_dict_items__register();
 py_Type pk_list__register();
diff --git a/include/pocketpy/objects/iterator.h b/include/pocketpy/objects/iterator.h
index c80be597..9e24dab0 100644
--- a/include/pocketpy/objects/iterator.h
+++ b/include/pocketpy/objects/iterator.h
@@ -13,3 +13,9 @@ typedef struct list_iterator {
     c11_vector* vec;
     int index;
 } list_iterator;
+
+typedef struct bytes_iterator {
+    unsigned char* data;
+    int length;
+    int index;
+} bytes_iterator;
diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h
index b012234a..077ecd0f 100644
--- a/include/pocketpy/pocketpy.h
+++ b/include/pocketpy/pocketpy.h
@@ -840,6 +840,7 @@ enum py_PredefinedType {
     tp_BaseException,
     tp_Exception,
     tp_bytes,
+    tp_bytes_iterator,
     tp_namedict,
     tp_locals,
     tp_code,
diff --git a/src/bindings/py_str.c b/src/bindings/py_str.c
index de5c556c..eda0dda2 100644
--- a/src/bindings/py_str.c
+++ b/src/bindings/py_str.c
@@ -796,6 +796,28 @@ static bool bytes__len__(int argc, py_Ref argv) {
     return true;
 }
 
+static bool bytes__iter__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    c11_bytes* self = py_touserdata(&argv[0]);
+    bytes_iterator* ud = py_newobject(py_retval(), tp_bytes_iterator, 1, sizeof(bytes_iterator));
+    ud->data = self->data;
+    ud->length = self->size;
+    ud->index = 0;
+    py_setslot(py_retval(), 0, argv);  // keep a reference to the bytes object
+    return true;
+}
+
+static bool bytes_iterator__next__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    bytes_iterator* ud = py_touserdata(argv);
+    if(ud->index < ud->length) {
+        py_newint(py_retval(), ud->data[ud->index]);
+        ud->index++;
+        return true;
+    }
+    return StopIteration();
+}
+
 py_Type pk_bytes__register() {
     py_Type type = pk_newtype("bytes", tp_object, NULL, NULL, false, true);
     // no need to dtor because the memory is controlled by the object
@@ -808,9 +830,17 @@ py_Type pk_bytes__register() {
     py_bindmagic(tp_bytes, __add__, bytes__add__);
     py_bindmagic(tp_bytes, __hash__, bytes__hash__);
     py_bindmagic(tp_bytes, __len__, bytes__len__);
+    py_bindmagic(tp_bytes, __iter__, bytes__iter__);
 
     py_bindmethod(tp_bytes, "decode", bytes_decode);
     return type;
 }
 
+py_Type pk_bytes_iterator__register() {
+    py_Type type = pk_newtype("bytes_iterator", tp_object, NULL, NULL, false, true);
+    py_bindmagic(type, __iter__, pk_wrapper__self);
+    py_bindmagic(type, __next__, bytes_iterator__next__);
+    return type;
+}
+
 #undef DEF_STR_CMP_OP
\ No newline at end of file
diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c
index f5573780..aa945e27 100644
--- a/src/interpreter/vm.c
+++ b/src/interpreter/vm.c
@@ -156,6 +156,7 @@ void VM__ctor(VM* self) {
     validate(tp_BaseException, pk_BaseException__register());
     validate(tp_Exception, pk_Exception__register());
     validate(tp_bytes, pk_bytes__register());
+    validate(tp_bytes_iterator, pk_bytes_iterator__register());
     validate(tp_namedict, pk_namedict__register());
     validate(tp_locals, pk_newtype("locals", tp_object, NULL, NULL, false, true));
     validate(tp_code, pk_code__register());