py/builtinimport.c: Add support for built-in packages.

jimmo · jimmo · commit 2a874dc7c47a · 2020-02-27T01:24:23.000+11:00
Allows a built-in module to add a sub-module as a member of its globals dict, and allow importing
it as "import foo.bar" (or any other variation on the import statement).

Note: This does give the slightly surprising behavior that "import foo" followed by "foo.bar.x"
will also work.

+56 bytes on PYBV11.
diff --git a/py/builtinimport.c b/py/builtinimport.c
@@ -354,6 +354,7 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
             DEBUG_printf("Previous path: =%.*s=\n", (int)vstr_len(&path), vstr_str(&path));
 
             mp_import_stat_t stat = MP_IMPORT_STAT_NO_EXIST;
+            bool store_on_parent = true;
 
             // Exact-match built-in (or already-loaded) takes priority.
             mp_obj_t module_obj = mp_module_get(mod_name);
@@ -377,7 +378,20 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
             } else {
                 DEBUG_printf("Searching for sub-module\n");
 
-                // TODO: packages for built-ins?
+                if (module_obj == MP_OBJ_NULL) {
+                    // See if the outer module has this module as an attr.
+                    // e.g. built-in modules can add modules to their globals dict
+                    // to provide something that behaves like built-in packages.
+                    mp_obj_t dest[2] = { MP_OBJ_NULL };
+                    mp_type_module.attr(outer_module_obj, qstr_from_strn(mod_str + last, i - last), dest);
+                    if (dest[0] != MP_OBJ_NULL) {
+                        // Ensure that it's actually a module.
+                        if (mp_obj_get_type(dest[0]) == &mp_type_module) {
+                            module_obj = dest[0];
+                            store_on_parent = false;
+                        }
+                    }
+                }
 
                 // Add the current part of the module name to the path.
                 // Again, need to do this even if already loaded (see find_file call above).
@@ -455,8 +469,8 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
                 }
             }
 
-            if (outer_module_obj != MP_OBJ_NULL) {
-                // If it's a sub-module, then make it available on
+            if (store_on_parent && outer_module_obj != MP_OBJ_NULL) {
+                // If it's a sub-module (not a built-in one), then make it available on
                 // the parent module.
                 qstr s = qstr_from_strn(mod_str + last, i - last);
                 mp_store_attr(outer_module_obj, s, module_obj);
