py/objfun: Implement function.__code__ and function constructor.

This allows retrieving the code object of a function using
`function.__code__`, and then reconstructing a function from a code object
using `FunctionType(code_object)`.

This feature is controlled by `MICROPY_PY_FUNCTION_ATTRS_CODE` and is
enabled at the full-features level.

Signed-off-by: Damien George <damien@micropython.org>

8 files changed, 133 insertions, 5 deletions

diff --git a/py/mpconfig.h b/py/mpconfig.h
index 76aff4681d..5c4d19bb56 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -1041,6 +1041,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_FUNCTION_ATTRS (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
 #endif
 
+// Whether to implement the __code__ attribute on functions, and function constructor
+#ifndef MICROPY_PY_FUNCTION_ATTRS_CODE
+#define MICROPY_PY_FUNCTION_ATTRS_CODE (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_FULL_FEATURES)
+#endif
+
 // Whether to support the descriptors __get__, __set__, __delete__
 // This costs some code size and makes load/store/delete of instance
 // attributes slower for the classes that use this feature
@@ -1135,7 +1140,7 @@ typedef double mp_float_t;
 #define MICROPY_PY_BUILTINS_CODE_BASIC      (2)
 #define MICROPY_PY_BUILTINS_CODE_FULL       (3)
 #ifndef MICROPY_PY_BUILTINS_CODE
-#define MICROPY_PY_BUILTINS_CODE            (MICROPY_PY_SYS_SETTRACE ? MICROPY_PY_BUILTINS_CODE_FULL : (MICROPY_PY_BUILTINS_COMPILE ? MICROPY_PY_BUILTINS_CODE_MINIMUM : MICROPY_PY_BUILTINS_CODE_NONE))
+#define MICROPY_PY_BUILTINS_CODE            (MICROPY_PY_SYS_SETTRACE ? MICROPY_PY_BUILTINS_CODE_FULL : (MICROPY_PY_FUNCTION_ATTRS_CODE ? MICROPY_PY_BUILTINS_CODE_BASIC : (MICROPY_PY_BUILTINS_COMPILE ? MICROPY_PY_BUILTINS_CODE_MINIMUM : MICROPY_PY_BUILTINS_CODE_NONE)))
 #endif
 
 // Whether to support dict.fromkeys() class method
diff --git a/py/objfun.c b/py/objfun.c
index 0b1b8c115f..a742c52672 100644
--- a/py/objfun.c
+++ b/py/objfun.c
@@ -28,6 +28,8 @@
 #include <string.h>
 #include <assert.h>
 
+#include "py/emitglue.h"
+#include "py/objcode.h"
 #include "py/objtuple.h"
 #include "py/objfun.h"
 #include "py/runtime.h"
@@ -151,6 +153,30 @@ qstr mp_obj_fun_get_name(mp_const_obj_t fun_in) {
     return name;
 }
 
+#if MICROPY_PY_FUNCTION_ATTRS_CODE
+static mp_obj_t fun_bc_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+    (void)type;
+    mp_arg_check_num(n_args, n_kw, 2, 2, false);
+
+    if (!mp_obj_is_type(args[0], &mp_type_code)) {
+        mp_raise_TypeError(NULL);
+    }
+    if (!mp_obj_is_type(args[1], &mp_type_dict)) {
+        mp_raise_TypeError(NULL);
+    }
+
+    mp_obj_code_t *code = MP_OBJ_TO_PTR(args[0]);
+    mp_obj_t globals = args[1];
+
+    mp_module_context_t *module_context = m_new_obj(mp_module_context_t);
+    module_context->module.base.type = &mp_type_module;
+    module_context->module.globals = MP_OBJ_TO_PTR(globals);
+    module_context->constants = *mp_code_get_constants(code);
+
+    return mp_make_function_from_proto_fun(mp_code_get_proto_fun(code), module_context, NULL);
+}
+#endif
+
 #if MICROPY_CPYTHON_COMPAT
 static void fun_bc_print(const mp_print_t *print, mp_obj_t o_in, mp_print_kind_t kind) {
     (void)kind;
@@ -340,9 +366,29 @@ void mp_obj_fun_bc_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
         mp_obj_fun_bc_t *self = MP_OBJ_TO_PTR(self_in);
         dest[0] = MP_OBJ_FROM_PTR(self->context->module.globals);
     }
+    #if MICROPY_PY_FUNCTION_ATTRS_CODE
+    if (attr == MP_QSTR___code__) {
+        const mp_obj_fun_bc_t *self = MP_OBJ_TO_PTR(self_in);
+        if ((self->base.type == &mp_type_fun_bc
+             || self->base.type == &mp_type_gen_wrap)
+            && self->child_table == NULL) {
+            #if MICROPY_PY_BUILTINS_CODE <= MICROPY_PY_BUILTINS_CODE_BASIC
+            dest[0] = mp_obj_new_code(self->context->constants, self->bytecode);
+            #else
+            dest[0] = mp_obj_new_code(self->context, self->rc, true);
+            #endif
+        }
+    }
+    #endif
 }
 #endif
 
+#if MICROPY_PY_FUNCTION_ATTRS_CODE
+#define FUN_BC_MAKE_NEW make_new, fun_bc_make_new,
+#else
+#define FUN_BC_MAKE_NEW
+#endif
+
 #if MICROPY_CPYTHON_COMPAT
 #define FUN_BC_TYPE_PRINT print, fun_bc_print,
 #else
@@ -359,6 +405,7 @@ MP_DEFINE_CONST_OBJ_TYPE(
     mp_type_fun_bc,
     MP_QSTR_function,
     MP_TYPE_FLAG_BINDS_SELF,
+    FUN_BC_MAKE_NEW
     FUN_BC_TYPE_PRINT
     FUN_BC_TYPE_ATTR
     call, fun_bc_call
diff --git a/tests/basics/fun_code.py b/tests/basics/fun_code.py
new file mode 100644
index 0000000000..59e1f7ec04
--- /dev/null
+++ b/tests/basics/fun_code.py
@@ -0,0 +1,36 @@
+# Test function.__code__ attribute.
+
+try:
+    (lambda: 0).__code__
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+
+def f():
+    return a
+
+
+ftype = type(f)
+
+# Test __code__ access and function constructor.
+code = f.__code__
+print(type(ftype(code, {})) is ftype)
+
+# Test instantiating multiple code's with different globals dicts.
+code = f.__code__
+f1 = ftype(code, {"a": 1})
+f2 = ftype(code, {"a": 2})
+print(f1(), f2())
+
+# Test bad first argument type.
+try:
+    ftype(None, {})
+except TypeError:
+    print("TypeError")
+
+# Test bad second argument type.
+try:
+    ftype(f.__code__, None)
+except TypeError:
+    print("TypeError")
diff --git a/tests/basics/fun_code_micropython.py b/tests/basics/fun_code_micropython.py
new file mode 100644
index 0000000000..2c319a2db8
--- /dev/null
+++ b/tests/basics/fun_code_micropython.py
@@ -0,0 +1,19 @@
+# Test MicroPython-specific restrictions of function.__code__ attribute.
+
+try:
+    (lambda: 0).__code__
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+
+def f_with_children():
+    def g():
+        pass
+
+
+# Can't access __code__ when function has children.
+try:
+    f_with_children.__code__
+except AttributeError:
+    print("AttributeError")
diff --git a/tests/basics/fun_code_micropython.py.exp b/tests/basics/fun_code_micropython.py.exp
new file mode 100644
index 0000000000..d169edffb4
--- /dev/null
+++ b/tests/basics/fun_code_micropython.py.exp
@@ -0,0 +1 @@
+AttributeError
diff --git a/tests/basics/subclass_native1.py b/tests/basics/subclass_native1.py
index 288a686d1a..74b377eac9 100644
--- a/tests/basics/subclass_native1.py
+++ b/tests/basics/subclass_native1.py
@@ -21,11 +21,9 @@ print(a + [20, 30, 40])
 # TODO: Faults
 #print(a + a)
 
-def foo():
-    print("hello from foo")
-
+# subclassing a type that doesn't have make_new at the C level (not allowed)
 try:
-    class myfunc(type(foo)):
+    class myfunc(type([].append)):
         pass
 except TypeError:
     print("TypeError")
diff --git a/tests/micropython/native_fun_attrs_code.py b/tests/micropython/native_fun_attrs_code.py
new file mode 100644
index 0000000000..d277a7b9b2
--- /dev/null
+++ b/tests/micropython/native_fun_attrs_code.py
@@ -0,0 +1,21 @@
+# Test MicroPython-specific restrictions of function.__code__ attribute.
+
+try:
+    (lambda: 0).__code__
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+import micropython
+
+
+@micropython.native
+def f_native():
+    pass
+
+
+# Can't access __code__ when function is native code.
+try:
+    f_native.__code__
+except AttributeError:
+    print("AttributeError")
diff --git a/tests/micropython/native_fun_attrs_code.py.exp b/tests/micropython/native_fun_attrs_code.py.exp
new file mode 100644
index 0000000000..d169edffb4
--- /dev/null
+++ b/tests/micropython/native_fun_attrs_code.py.exp
@@ -0,0 +1 @@
+AttributeError