fix super()

include/typings/pkpy.pyi

@@ -7,7 +7,7 @@ class TValue[T]:
     @property
     def value(self) -> T: ...
 
-TValue_int = TValue[int]
-TValue_float = TValue[float]
-TValue_vec2i = TValue[vec2i]
-TValue_vec2 = TValue[vec2]
+# TValue_int = TValue[int]
+# TValue_float = TValue[float]
+# TValue_vec2i = TValue[vec2i]
+# TValue_vec2 = TValue[vec2]

src/interpreter/ceval.c

@@ -974,8 +974,13 @@ FrameResult VM__run_top_frame(VM* self) {
             case OP_STORE_CLASS_ATTR: {
                 assert(self->__curr_class);
                 py_Name name = byte.arg;
-                if(py_istype(TOP(), tp_function)) {
-                    Function* ud = py_touserdata(TOP());
+                // TOP() can be a function, classmethod or custom decorator
+                py_Ref actual_func = TOP();
+                if(actual_func->type == tp_classmethod) {
+                    actual_func = py_getslot(actual_func, 0);
+                }
+                if(actual_func->type == tp_function) {
+                    Function* ud = py_touserdata(actual_func);
                     ud->clazz = self->__curr_class->_obj;
                 }
                 py_setdict(self->__curr_class, name, TOP());

src/public/internal.c

@@ -167,9 +167,6 @@ bool pk_loadmethod(py_StackRef self, py_Name name) {
         if(py_istype(self, tp_type)) {
             // T.__new__(...)
             type = py_totype(self);
-        } else if(py_istype(self, tp_super)) {
-            // super().__new__(...)
-            type = *(py_Type*)py_touserdata(self);
         } else {
             // invalid usage of `__new__`
             return false;

src/public/modules.c

@@ -756,33 +756,47 @@ py_Type pk_nativefunc__register() {
 }
 
 static bool super__new__(int argc, py_Ref argv) {
-    py_Type* class_arg = py_newobject(py_retval(), tp_super, 1, sizeof(py_Type));
+    py_Type class_arg = 0;
     Frame* frame = pk_current_vm->top_frame;
-    *class_arg = 0;
     py_Ref self_arg = NULL;
     if(argc == 1) {
         // super()
         if(frame->has_function) {
-            Function* func = py_touserdata(frame->p0);
-            *class_arg = *(py_Type*)PyObject__userdata(func->clazz);
+            py_TValue* callable = frame->p0;
+            if(callable->type == tp_boundmethod) callable = py_getslot(frame->p0, 1);
+            if(callable->type == tp_function) {
+                Function* func = py_touserdata(callable);
+                if(func->clazz != NULL) {
+                    class_arg = *(py_Type*)PyObject__userdata(func->clazz);
                     if(frame->co->nlocals > 0) self_arg = &frame->locals[0];
                 }
+            }
+        }
         if(class_arg == 0 || self_arg == NULL) return RuntimeError("super(): no arguments");
+        if(self_arg->type == tp_type) {
+            // f(cls, ...)
+            class_arg = pk__type_info(class_arg)->base;
+            if(class_arg == 0) return RuntimeError("super(): base class is invalid");
+            py_assign(py_retval(), py_tpobject(class_arg));
+            return true;
+        }
     } else if(argc == 3) {
-        // super(type, obj)
+        // super(type[T], obj)
         PY_CHECK_ARG_TYPE(1, tp_type);
-        *class_arg = py_totype(py_arg(1));
+        class_arg = py_totype(py_arg(1));
         self_arg = py_arg(2);
-        if(!py_isinstance(self_arg, *class_arg)) {
+        if(!py_isinstance(self_arg, class_arg)) {
             return TypeError("super(type, obj): obj must be an instance of type");
         }
     } else {
         return TypeError("super() takes 0 or 2 arguments");
     }
 
-    *class_arg = pk__type_info(*class_arg)->base;
-    if(*class_arg == 0) return RuntimeError("super(): base class is invalid");
+    class_arg = pk__type_info(class_arg)->base;
+    if(class_arg == 0) return RuntimeError("super(): base class is invalid");
 
+    py_Type* p_class_arg = py_newobject(py_retval(), tp_super, 1, sizeof(py_Type));
+    *p_class_arg = class_arg;
     py_setslot(py_retval(), 0, self_arg);
     return true;
 }

src/public/py_ops.c

@@ -87,12 +87,6 @@ int py_next(py_Ref val) {
 bool py_getattr(py_Ref self, py_Name name) {
     // https://docs.python.org/3/howto/descriptor.html#invocation-from-an-instance
     py_Type type = self->type;
-    // handle super() proxy (disabled)
-    // if(py_istype(self, tp_super)) {
-    //     self = py_getslot(self, 0);
-    //     type = *(py_Type*)py_touserdata(self);
-    // }
-
     py_Ref cls_var = py_tpfindname(type, name);
     if(cls_var) {
         // handle descriptor
@@ -183,12 +177,6 @@ bool py_getattr(py_Ref self, py_Name name) {
 
 bool py_setattr(py_Ref self, py_Name name, py_Ref val) {
     py_Type type = self->type;
-    // handle super() proxy (disabled)
-    // if(py_istype(self, tp_super)) {
-    //     self = py_getslot(self, 0);
-    //     type = *(py_Type*)py_touserdata(self);
-    // }
-
     py_Ref cls_var = py_tpfindname(type, name);
     if(cls_var) {
         // handle descriptor

tests/41_class_ex.py

@@ -154,3 +154,17 @@ assert B.static_method(123) == 123
 assert B.class_method('123') == 'B123'
 assert B().static_method(123) == 123
 assert B().class_method('123') == 'B123'
+
+# test super() with classmethod
+class BaseClass:
+    @classmethod
+    def f(cls):
+        return 'BaseClass'
+
+class DerivedClass(BaseClass):
+    @classmethod
+    def f(cls):
+        return super().f()
+
+    
+assert DerivedClass.f() == 'BaseClass'

tests/99_extras.py

@@ -78,6 +78,7 @@ class Z:
 
 class B(Z):
     def __new__(cls, x):
+        assert super() is Z
         return super().__new__(cls, x)
 
 assert Z(1) == (Z, 1)
@@ -87,6 +88,7 @@ from pkpy import TValue
 
 class fixed(TValue[int]):
     def __new__(cls, value: str):
+        assert super() is TValue[int]
         return super().__new__(cls, int(value))
     
 assert fixed('123').value == 123