diff --git a/src/pocketpy.h b/src/pocketpy.h index e3325176..ed922957 100644 --- a/src/pocketpy.h +++ b/src/pocketpy.h @@ -223,6 +223,7 @@ inline void init_builtins(VM* _vm) { _vm->bind_method<0>("range", "__iter__", CPP_LAMBDA( vm->PyIter(RangeIter(vm, args[0])) )); + _vm->_type_info("range")->m__iter__ = PyGetIter; _vm->bind_method<0>("NoneType", "__repr__", CPP_LAMBDA(VAR("None"))); _vm->bind_method<0>("NoneType", "__json__", CPP_LAMBDA(VAR("null"))); @@ -365,6 +366,7 @@ inline void init_builtins(VM* _vm) { _vm->bind_method<0>("str", "__str__", CPP_LAMBDA(args[0])); _vm->bind_method<0>("str", "__iter__", CPP_LAMBDA(vm->PyIter(StringIter(vm, args[0])))); + _vm->_type_info("str")->m__iter__ = PyGetIter; _vm->bind_method<0>("str", "__repr__", [](VM* vm, ArgsView args) { const Str& self = _CAST(Str&, args[0]); @@ -539,13 +541,13 @@ inline void init_builtins(VM* _vm) { _vm->bind_method<0>("list", "__iter__", [](VM* vm, ArgsView args) { return vm->PyIter(ArrayIter(vm, args[0])); }); - _vm->bind_method<1>("list", "__getitem__", [](VM* vm, ArgsView args) { return PyArrayGetItem(vm, args[0], args[1]); }); _vm->bind_method<2>("list", "__setitem__", [](VM* vm, ArgsView args) { return PyListSetItem(vm, args[0], args[1], args[2]); }); + _vm->_type_info("list")->m__iter__ = PyGetIter>; _vm->_type_info("list")->m__getitem__ = PyArrayGetItem; _vm->_type_info("list")->m__setitem__ = PyListSetItem; @@ -566,6 +568,7 @@ inline void init_builtins(VM* _vm) { _vm->bind_method<0>("tuple", "__iter__", [](VM* vm, ArgsView args) { return vm->PyIter(ArrayIter(vm, args[0])); }); + _vm->_type_info("tuple")->m__iter__ = PyGetIter>; _vm->bind_method<1>("tuple", "__getitem__", [](VM* vm, ArgsView args) { return PyArrayGetItem(vm, args[0], args[1]); diff --git a/src/vm.h b/src/vm.h index 8b173e86..9767f0ac 100644 --- a/src/vm.h +++ b/src/vm.h @@ -68,6 +68,7 @@ struct PyTypeInfo{ // cached special methods PyObject* (*m__getitem__)(VM* vm, PyObject*, PyObject*) = nullptr; PyObject* (*m__setitem__)(VM* vm, PyObject*, PyObject*, PyObject*) = nullptr; + PyObject* (*m__iter__)(VM* vm, PyObject*) = nullptr; }; struct FrameId{ @@ -136,6 +137,8 @@ public: PyObject* asIter(PyObject* obj){ if(is_type(obj, tp_iterator)) return obj; + const PyTypeInfo* ti = _inst_type_info(obj); + if(ti->m__iter__) return ti->m__iter__(this, obj); PyObject* self; PyObject* iter_f = get_unbound_method(obj, __iter__, &self, false); if(self != PY_NULL) return call_method(self, iter_f); @@ -1278,4 +1281,9 @@ inline PyObject* PyStrGetItem(VM* vm, PyObject* obj, PyObject* index){ return VAR(self.u8_getitem(i)); } +template +inline PyObject* PyGetIter(VM* vm, PyObject* obj){ + return vm->PyIter(T(vm, obj)); +} + } // namespace pkpy \ No newline at end of file