add enum module

include/pocketpy/common/_generated.h

@@ -3,7 +3,6 @@
 
 const char* load_kPythonLib(const char* name);
 
-extern const char kPythonLibs__enum[];
 extern const char kPythonLibs_bisect[];
 extern const char kPythonLibs_builtins[];
 extern const char kPythonLibs_cmath[];

include/pocketpy/interpreter/modules.h

@@ -10,4 +10,5 @@ void pk__add_module_json();
 void pk__add_module_gc();
 void pk__add_module_time();
 void pk__add_module_easing();
-void pk__add_module_traceback();
+void pk__add_module_traceback();
+void pk__add_module_enum();

include/pocketpy/interpreter/vm.h

@@ -48,7 +48,6 @@ typedef struct VM {
     py_TValue last_retval;
     py_TValue curr_exception;
     bool is_curr_exc_handled;  // handled by try-except block but not cleared yet
-    bool is_stopiteration;
 
     py_TValue reg[8];  // users' registers
 

include/pocketpy/pocketpy.h

@@ -295,6 +295,10 @@ PK_EXPORT void py_setdict(py_Ref self, py_Name name, py_Ref val);
 PK_EXPORT bool py_deldict(py_Ref self, py_Name name);
 /// Prepare an insertion to the object's `__dict__`.
 PK_EXPORT py_ItemRef py_emplacedict(py_Ref self, py_Name name);
+/// Apply a function to all items in the object's `__dict__`.
+/// Return `true` if the function is successful for all items.
+/// NOTE: Be careful if `f` modifies the object's `__dict__`.
+PK_EXPORT bool py_applydict(py_Ref self, bool (*f)(py_Name name, py_Ref val, void* ctx), void* ctx) PY_RAISE;
 
 /// Get the i-th slot of the object.
 /// The object must have slots and `i` must be in valid range.

python/_enum.py

@@ -1,11 +0,0 @@
-class Enum:
-    def __init__(self, name, value):
-        self.name = name
-        self.value = value
-
-    def __str__(self):
-        return f'{type(self).__name__}.{self.name}'
-    
-    def __repr__(self):
-        return f'<{str(self)}: {self.value!r}>'
-    

src/interpreter/vm.c

@@ -68,7 +68,6 @@ void VM__ctor(VM* self) {
     self->last_retval = *py_NIL;
     self->curr_exception = *py_NIL;
     self->is_curr_exc_handled = false;
-    self->is_stopiteration = false;
 
     self->__curr_class = NULL;
 
@@ -206,6 +205,7 @@ void VM__ctor(VM* self) {
     pk__add_module_time();
     pk__add_module_easing();
     pk__add_module_traceback();
+    pk__add_module_enum();
 
     // add python builtins
     do {

src/modules/enum.c

@@ -0,0 +1,91 @@
+#include "pocketpy/pocketpy.h"
+#include "pocketpy/interpreter/vm.h"
+#include "pocketpy/common/sstream.h"
+
+static bool Enum__wrapper_field(py_Name name, py_Ref value, void* ctx) {
+    c11_sv name_sv = py_name2sv(name);
+    if(name_sv.size == 0 || name_sv.data[0] == '_') return true;
+    py_push(ctx);
+    py_pushnil();
+    py_newstr(py_pushtmp(), py_name2str(name));
+    py_push(value);
+    bool ok = py_vectorcall(2, 0);
+    if(!ok) return false;
+    py_assign(value, py_retval());
+    return true;
+}
+
+static void Enum__on_end_subclass(py_TypeInfo* derived_ti) {
+    derived_ti->is_sealed = true;
+    py_applydict(&derived_ti->self, Enum__wrapper_field, &derived_ti->self);
+}
+
+static bool Enum__new__(int argc, py_Ref argv) {
+    py_newobject(py_retval(), py_totype(argv), 2, 0);
+    return true;
+}
+
+static bool Enum__init__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(3);
+    PY_CHECK_ARG_TYPE(1, tp_str);
+    py_setslot(argv, 0, py_arg(1));
+    py_setslot(argv, 1, py_arg(2));
+    py_newnone(py_retval());
+    return true;
+}
+
+static bool Enum__str__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    // f'{type(self).__name__}.{self.name}'
+    c11_sbuf buf;
+    c11_sbuf__ctor(&buf);
+    c11_sbuf__write_cstr(&buf, py_tpname(argv->type));
+    c11_sbuf__write_char(&buf, '.');
+    c11_sbuf__write_cstr(&buf, py_tostr(py_getslot(argv, 0)));
+    c11_sbuf__py_submit(&buf, py_retval());
+    return true;
+}
+
+static bool Enum__repr__(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    // f'<{str(self)}: {self.value!r}>'
+    if(!py_str(argv)) return false;
+    py_push(py_retval());  // str(self)
+    if(!py_repr(py_getslot(argv, 1))) return false;
+    py_push(py_retval());  // repr(self.value)
+    c11_sbuf buf;
+    c11_sbuf__ctor(&buf);
+    c11_sbuf__write_cstr(&buf, "<");
+    c11_sbuf__write_cstr(&buf, py_tostr(py_peek(-2)));
+    c11_sbuf__write_cstr(&buf, ": ");
+    c11_sbuf__write_cstr(&buf, py_tostr(py_peek(-1)));
+    c11_sbuf__write_cstr(&buf, ">");
+    c11_sbuf__py_submit(&buf, py_retval());
+    py_shrink(2);
+    return true;
+}
+
+static bool Enum__name(int argc, py_Ref argv) {
+    py_assign(py_retval(), py_getslot(argv, 0));
+    return true;
+}
+
+static bool Enum__value(int argc, py_Ref argv) {
+    py_assign(py_retval(), py_getslot(argv, 1));
+    return true;
+}
+
+void pk__add_module_enum() {
+    py_Ref mod = py_newmodule("enum");
+    py_Type type = py_newtype("Enum", tp_object, mod, NULL);
+
+    py_bindmagic(type, __new__, Enum__new__);
+    py_bindmagic(type, __init__, Enum__init__);
+    py_bindmagic(type, __str__, Enum__str__);
+    py_bindmagic(type, __repr__, Enum__repr__);
+    py_bindproperty(type, "name", Enum__name, NULL);
+    py_bindproperty(type, "value", Enum__value, NULL);
+
+    py_TypeInfo* ti = c11__at(py_TypeInfo, &pk_current_vm->types, type);
+    ti->on_end_subclass = Enum__on_end_subclass;
+}

src/public/internal.c

@@ -105,7 +105,12 @@ bool py_callcfunc(py_CFunction f, int argc, py_Ref argv) {
     py_StackRef p0 = py_peek(0);
     py_newnil(py_retval());
     bool ok = f(argc, argv);
-    if(!ok) return false;
+    if(!ok) {
+        if(!py_checkexc(true)) {
+            c11__abort("py_CFunction returns `false` but no exception is set!");
+        }
+        return false;
+    }
     if(py_peek(0) != p0) {
         c11__abort("py_CFunction corrupts the stack! Did you forget to call `py_pop()`?");
     }
@@ -219,9 +224,4 @@ bool pk_callmagic(py_Name name, int argc, py_Ref argv) {
     return py_call(tmp, argc, argv);
 }
 
-bool StopIteration() {
-    VM* vm = pk_current_vm;
-    assert(!vm->is_stopiteration);  // flag is already set
-    vm->is_stopiteration = true;
-    return false;
-}
+bool StopIteration() { return py_exception(tp_StopIteration, ""); }

src/public/py_exception.c

@@ -152,7 +152,6 @@ void py_clearexc(py_StackRef p0) {
     vm->last_retval = *py_NIL;
     vm->curr_exception = *py_NIL;
     vm->is_curr_exc_handled = false;
-    vm->is_stopiteration = false;
     vm->__curr_class = NULL;
     if(p0) vm->stack.sp = p0;
 }

src/public/py_ops.c

@@ -70,20 +70,17 @@ bool py_iter(py_Ref val) {
 
 int py_next(py_Ref val) {
     VM* vm = pk_current_vm;
-    vm->is_stopiteration = false;
     py_Ref tmp = py_tpfindmagic(val->type, __next__);
     if(!tmp) {
         TypeError("'%t' object is not an iterator", val->type);
         return -1;
     }
-    if(py_call(tmp, 1, val)) return true;
+    if(py_call(tmp, 1, val)) return 1;
     if(vm->curr_exception.type == tp_StopIteration) {
         py_clearexc(NULL);
-        vm->is_stopiteration = true;
+        return 0;
     }
-    int retval = vm->is_stopiteration ? 0 : -1;
-    vm->is_stopiteration = false;
-    return retval;
+    return -1;
 }
 
 bool py_getattr(py_Ref self, py_Name name) {

src/public/stack_ops.c

@@ -34,6 +34,17 @@ py_ItemRef py_emplacedict(py_Ref self, py_Name name){
     return py_getdict(self, name);
 }
 
+bool py_applydict(py_Ref self, bool (*f)(py_Name, py_Ref, void *), void *ctx){
+    assert(self && self->is_ptr);
+    NameDict* dict = PyObject__dict(self->_obj);
+    for(int i = 0; i < dict->length; i++){
+        NameDict_KV* kv = c11__at(NameDict_KV, dict, i);
+        bool ok = f(kv->key, &kv->value, ctx);
+        if(!ok) return false;
+    }
+    return true;
+}
+
 bool py_deldict(py_Ref self, py_Name name) {
     assert(self && self->is_ptr);
     if(!py_ismagicname(name) || self->type != tp_type) {

tests/82_enum.py

@@ -1,5 +1,3 @@
-exit()
-
 from enum import Enum
 
 class A(Enum):