add enum module

docs/modules/enum.md

@@ -0,0 +1,24 @@
+---
+icon: package
+label: enum
+---
+
+### `enum.Enum`
+
+Base class for creating enumerated constants.
+
+Example:
+
+```python
+from enum import Enum
+
+class Color(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+print(Color.RED)    # Color.RED
+print(Color.RED.name)    # 'RED'
+print(Color.RED.value)    # 1
+```
+

include/pocketpy/modules.h

@@ -11,5 +11,6 @@ void add_module_traceback(VM* vm);
 void add_module_dis(VM* vm);
 void add_module_gc(VM* vm);
 void add_module_line_profiler(VM* vm);
+void add_module_enum(VM* vm);
 
 }   // namespace pkpy

include/pocketpy/vm.h

@@ -82,6 +82,9 @@ struct PyTypeInfo{
     void (*m__setattr__)(VM* vm, PyObject*, StrName, PyObject*) = nullptr;
     PyObject* (*m__getattr__)(VM* vm, PyObject*, StrName) = nullptr;
     bool (*m__delattr__)(VM* vm, PyObject*, StrName) = nullptr;
+
+    // backdoors
+    void (*on_end_subclass)(VM* vm, PyTypeInfo*) = nullptr;
 };
 
 typedef void(*PrintFunc)(const char*, int);

python/_enum.py

@@ -0,0 +1,11 @@
+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/ceval.cpp

@@ -814,6 +814,12 @@ __NEXT_STEP:;
         PK_ASSERT(_curr_class != nullptr);
         StrName _name(byte.arg);
         frame->_module->attr().set(_name, _curr_class);
+        // call on_end_subclass
+        PyTypeInfo* ti = &_all_types[PK_OBJ_GET(Type, _curr_class)];
+        if(ti->base != tp_object){
+            PyTypeInfo* base_ti = &_all_types[ti->base];
+            if(base_ti->on_end_subclass) base_ti->on_end_subclass(this, ti);
+        }
         _curr_class = nullptr;
     } DISPATCH();
     TARGET(STORE_CLASS_ATTR){

src/modules.cpp

@@ -315,4 +315,23 @@ void add_module_line_profiler(VM *vm){
     LineProfilerW::register_class(vm, mod);
 }
 
+
+void add_module_enum(VM* vm){
+    PyObject* mod = vm->new_module("enum");
+    CodeObject_ code = vm->compile(kPythonLibs__enum, "enum.py", EXEC_MODE);
+    vm->_exec(code, mod);
+    PyObject* Enum = mod->attr("Enum");
+    vm->_all_types[PK_OBJ_GET(Type, Enum).index].on_end_subclass = \
+        [](VM* vm, PyTypeInfo* new_ti){
+            new_ti->subclass_enabled = false;    // Enum class cannot be subclassed twice
+            NameDict& attr = new_ti->obj->attr();
+            for(auto [k, v]: attr.items()){
+                // wrap every attribute
+                std::string_view k_sv = k.sv();
+                if(k_sv.empty() || k_sv[0] == '_') continue;
+                attr.set(k, vm->call(new_ti->obj, VAR(k_sv), v));
+            }
+        };
+}
+
 }   // namespace pkpy

src/pocketpy.cpp

@@ -1618,6 +1618,7 @@ void VM::post_init(){
     add_module_collections(this);
     add_module_array2d(this);
     add_module_line_profiler(this);
+    add_module_enum(this);
 
 #ifdef PK_USE_CJSON
     add_module_cjson(this);

tests/85_enum.py

@@ -0,0 +1,28 @@
+from enum import Enum
+
+class A(Enum):
+    a = 1
+    b = '2'
+    c = None
+
+assert str(A.a) == 'A.a'
+assert repr(A.a) == '<A.a: 1>'
+
+assert str(A.b) == 'A.b'
+assert repr(A.b) == "<A.b: '2'>"
+
+assert str(A.c) == 'A.c'
+assert repr(A.c) == '<A.c: None>'
+
+assert A.a == A.a
+assert A.a != A.b
+assert A.a != A.c
+
+assert A.a.name == 'a'
+assert A.a.value == 1
+
+assert A.b.name == 'b'
+assert A.b.value == '2'
+
+assert A.c.name == 'c'
+assert A.c.value is None