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());