py: Improve __build_class__.

5 files changed, 54 insertions, 12 deletions

diff --git a/py/builtin.c b/py/builtin.c
index 6babc76692..4ee5c57c62 100644
--- a/py/builtin.c
+++ b/py/builtin.c
@@ -15,22 +15,59 @@
 #include "map.h"
 #include "builtin.h"
 
-mp_obj_t mp_builtin___build_class__(mp_obj_t o_class_fun, mp_obj_t o_class_name) {
+// args[0] is function from class body
+// args[1] is class name
+// args[2:] are base objects
+mp_obj_t mp_builtin___build_class__(int n_args, const mp_obj_t *args) {
+    assert(2 <= n_args);
+
     // we differ from CPython: we set the new __locals__ object here
     mp_map_t *old_locals = rt_locals_get();
     mp_map_t *class_locals = mp_map_new(MP_MAP_QSTR, 0);
     rt_locals_set(class_locals);
 
     // call the class code
-    rt_call_function_1(o_class_fun, (mp_obj_t)0xdeadbeef);
+    mp_obj_t cell = rt_call_function_1(args[0], (mp_obj_t)0xdeadbeef);
 
     // restore old __locals__ object
     rt_locals_set(old_locals);
 
-    // create and return the new class
-    return mp_obj_new_class(class_locals);
+    /*
+    // get the class type (meta object) from the base objects
+    mp_obj_t meta;
+    if (n_args == 2) {
+        // no explicit bases, so use 'type'
+        meta = (mp_obj_t)&mp_const_type;
+    } else {
+        // use type of first base object
+        meta = mp_obj_get_type(args[2]);
+    }
+    */
+
+    // TODO do proper metaclass resolution for multiple base objects
+
+    /*
+    // create the new class using a call to the meta object
+    // (arguments must be backwards in the array)
+    mp_obj_t meta_args[3];
+    meta_args[2] = args[1]; // class name
+    meta_args[1] = mp_obj_new_tuple(n_args - 2, args + 2); // tuple of bases
+    meta_args[0] = class_locals; // dict of members TODO, currently is a map
+    mp_obj_t new_class = rt_call_function_n(meta, 3, meta_args);
+    */
+    // create the new class
+    mp_obj_t new_class = mp_obj_new_class(class_locals);
+
+    // store into cell if neede
+    if (cell != mp_const_none) {
+        mp_obj_cell_set(cell, new_class);
+    }
+
+    return new_class;
 }
 
+MP_DEFINE_CONST_FUN_OBJ_VAR(mp_builtin___build_class___obj, 2, mp_builtin___build_class__);
+
 mp_obj_t mp_builtin___repl_print__(mp_obj_t o) {
     if (o != mp_const_none) {
         mp_obj_print(o);
@@ -281,12 +318,7 @@ mp_obj_t mp_builtin_sum(int n_args, const mp_obj_t *args) {
 
 static mp_obj_t mp_builtin_type(mp_obj_t o_in) {
     // TODO implement the 3 argument version of type()
-    if (MP_OBJ_IS_SMALL_INT(o_in)) {
-        return (mp_obj_t)&int_type;
-    } else {
-        mp_obj_base_t *o = o_in;
-        return (mp_obj_t)o->type;
-    }
+    return mp_obj_get_type(o_in);
 }
 
 MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_type_obj, mp_builtin_type);
diff --git a/py/builtin.h b/py/builtin.h
index 9806624206..2cd0ef15f2 100644
--- a/py/builtin.h
+++ b/py/builtin.h
@@ -1,6 +1,6 @@
 // TODO convert all these to objects using MP_DECLARE and MP_DEFINE
 
-mp_obj_t mp_builtin___build_class__(mp_obj_t o_class_fun, mp_obj_t o_class_name);
+MP_DECLARE_CONST_FUN_OBJ(mp_builtin___build_class___obj);
 mp_obj_t mp_builtin___import__(int n, mp_obj_t *args);
 mp_obj_t mp_builtin___repl_print__(mp_obj_t o);
 mp_obj_t mp_builtin_abs(mp_obj_t o_in);
diff --git a/py/obj.c b/py/obj.c
index 77580e1fee..7d0122c8fe 100644
--- a/py/obj.c
+++ b/py/obj.c
@@ -13,6 +13,15 @@
 #include "runtime.h"
 #include "map.h"
 
+mp_obj_t mp_obj_get_type(mp_obj_t o_in) {
+    if (MP_OBJ_IS_SMALL_INT(o_in)) {
+        return (mp_obj_t)&int_type;
+    } else {
+        mp_obj_base_t *o = o_in;
+        return (mp_obj_t)o->type;
+    }
+}
+
 const char *mp_obj_get_type_str(mp_obj_t o_in) {
     if (MP_OBJ_IS_SMALL_INT(o_in)) {
         return "int";
diff --git a/py/obj.h b/py/obj.h
index 71b43cb8d1..208b23442f 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -161,6 +161,7 @@ mp_obj_t mp_obj_new_class(struct _mp_map_t *class_locals);
 mp_obj_t mp_obj_new_instance(mp_obj_t clas);
 mp_obj_t mp_obj_new_module(qstr module_name);
 
+mp_obj_t mp_obj_get_type(mp_obj_t o_in);
 const char *mp_obj_get_type_str(mp_obj_t o_in);
 
 void mp_obj_print_helper(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t o_in);
diff --git a/py/runtime.c b/py/runtime.c
index 5c476d9ecf..c0ed3b1fac 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -90,7 +90,7 @@ void rt_init(void) {
     mp_qstr_map_lookup(&map_builtins, MP_QSTR_Ellipsis, true)->value = mp_const_ellipsis;
 
     // built-in core functions
-    mp_qstr_map_lookup(&map_builtins, MP_QSTR___build_class__, true)->value = rt_make_function_2(mp_builtin___build_class__);
+    mp_qstr_map_lookup(&map_builtins, MP_QSTR___build_class__, true)->value = (mp_obj_t)&mp_builtin___build_class___obj;
     mp_qstr_map_lookup(&map_builtins, MP_QSTR___repl_print__, true)->value = rt_make_function_1(mp_builtin___repl_print__);
 
     // built-in types