extmod/modmarshal: Support marshal.dumps of functions with children.

This commit adds support to the `marshal` module to be able to dump
functions that have child functions.  For example:

    import marshal

    def f():
        def child():
            return 1
        return child

    marshal.dumps(f.__code__)

It also covers the case of marshalling functions that use list
comprehensions, because a list comprehension uses a child function.

This is made possible by the newly enhanced
`mp_raw_code_save_fun_to_bytes()` that can now handle nested functions.

Unmarshalling via `marshal.loads()` already supports nested functions
because it uses the standard `mp_raw_code_load_mem()` function which is
used to import mpy files (and hence can handle all possibilities).

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

extmod/modmarshal.c

@@ -36,17 +36,7 @@ static mp_obj_t marshal_dumps(mp_obj_t value_in) {
     if (mp_obj_is_type(value_in, &mp_type_code)) {
         mp_obj_code_t *code = MP_OBJ_TO_PTR(value_in);
         const void *proto_fun = mp_code_get_proto_fun(code);
-        const uint8_t *bytecode;
-        if (mp_proto_fun_is_bytecode(proto_fun)) {
-            bytecode = proto_fun;
-        } else {
-            const mp_raw_code_t *rc = proto_fun;
-            if (!(rc->kind == MP_CODE_BYTECODE && rc->children == NULL)) {
-                mp_raise_ValueError(MP_ERROR_TEXT("function must be bytecode with no children"));
-            }
-            bytecode = rc->fun_data;
-        }
-        return mp_raw_code_save_fun_to_bytes(mp_code_get_constants(code), bytecode);
+        return mp_raw_code_save_fun_to_bytes(mp_code_get_constants(code), proto_fun);
     } else {
         mp_raise_ValueError(MP_ERROR_TEXT("unmarshallable object"));
     }

tests/extmod/marshal_fun_nested.py

@@ -0,0 +1,79 @@
+# Test the marshal module, with functions that have children.
+
+try:
+    import marshal
+
+    (lambda: 0).__code__
+except (AttributeError, ImportError):
+    print("SKIP")
+    raise SystemExit
+
+
+def f_with_child():
+    def child():
+        return a
+
+    return child
+
+
+def f_with_child_defargs():
+    def child(a="default"):
+        return a
+
+    return child
+
+
+def f_with_child_closure():
+    a = "closure 1"
+
+    def child():
+        return a
+
+    a = "closure 2"
+    return child
+
+
+def f_with_child_closure_defargs():
+    a = "closure defargs 1"
+
+    def child(b="defargs default"):
+        return (a, b)
+
+    a = "closure defargs 1"
+    return child
+
+
+def f_with_list_comprehension(a):
+    return [i + a for i in range(4)]
+
+
+ftype = type(lambda: 0)
+
+# Test function with a child.
+f = ftype(marshal.loads(marshal.dumps(f_with_child.__code__)), {"a": "global"})
+print(f()())
+
+# Test function with a child that has default arguments.
+f = ftype(marshal.loads(marshal.dumps(f_with_child_defargs.__code__)), {})
+print(f()())
+print(f()("non-default"))
+
+# Test function with a child that is a closure.
+f = ftype(marshal.loads(marshal.dumps(f_with_child_closure.__code__)), {})
+print(f()())
+
+# Test function with a child that is a closure and has default arguments.
+f = ftype(marshal.loads(marshal.dumps(f_with_child_closure_defargs.__code__)), {})
+print(f()())
+print(f()("defargs non-default"))
+
+# Test function with a list comprehension (which will be an anonymous child).
+f = ftype(marshal.loads(marshal.dumps(f_with_list_comprehension.__code__)), {})
+print(f(10))
+
+# Test child within a module (the outer scope).
+code = compile("def child(a): return a", "", "exec")
+f = marshal.loads(marshal.dumps(code))
+ctx = {}
+exec(f, ctx)
+print(ctx["child"]("arg"))

tests/extmod/marshal_micropython.py

@@ -1,21 +0,0 @@
-# Test the marshal module, MicroPython-specific functionality.
-
-try:
-    import marshal
-except ImportError:
-    print("SKIP")
-    raise SystemExit
-
-import unittest
-
-
-class Test(unittest.TestCase):
-    def test_function_with_children(self):
-        # Can't marshal a function with children (in this case the module has a child function f).
-        code = compile("def f(): pass", "", "exec")
-        with self.assertRaises(ValueError):
-            marshal.dumps(code)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/micropython/native_marshal.py

@@ -0,0 +1,45 @@
+# Test the marshal module in combination with native/viper functions.
+
+try:
+    import marshal
+
+    (lambda: 0).__code__
+except (AttributeError, ImportError):
+    print("SKIP")
+    raise SystemExit
+
+import unittest
+
+
+def f_native():
+    @micropython.native
+    def g():
+        pass
+
+    return g
+
+
+def f_viper():
+    @micropython.viper
+    def g():
+        pass
+
+    return g
+
+
+class Test(unittest.TestCase):
+    def test_native_function(self):
+        # Can't marshal a function with native code.
+        code = f_native.__code__
+        with self.assertRaises(ValueError):
+            marshal.dumps(code)
+
+    def test_viper_function(self):
+        # Can't marshal a function with viper code.
+        code = f_viper.__code__
+        with self.assertRaises(ValueError):
+            marshal.dumps(code)
+
+
+if __name__ == "__main__":
+    unittest.main()