Optimize import mechanism

ports/unix/main.c

@@ -629,7 +629,9 @@ MP_NOINLINE int main_(int argc, char **argv) {
                     return handle_uncaught_exception(nlr.ret_val) & 0xff;
                 }
 
-                if (mp_obj_is_package(mod) && !subpkg_tried) {
+                mp_obj_t dest[2];
+                mp_load_method_maybe(mod, MP_QSTR___path__, dest);
+                if (dest[0] != MP_OBJ_NULL && !subpkg_tried) {
                     subpkg_tried = true;
                     vstr_t vstr;
                     int len = strlen(argv[a + 1]);

py/builtin.h

@@ -64,7 +64,13 @@ MP_DECLARE_CONST_FUN_OBJ_KW(mp_builtin_open_obj);
 
 #endif
 
+// A port can provide its own import handler by defining mp_builtin___import__.
+#ifndef mp_builtin___import__
+#define mp_builtin___import__ mp_builtin___import___default
+#endif
 mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args);
+mp_obj_t mp_builtin___import___default(size_t n_args, const mp_obj_t *args);
+
 mp_obj_t mp_micropython_mem_info(size_t n_args, const mp_obj_t *args);
 
 MP_DECLARE_CONST_FUN_OBJ_VAR(mp_builtin___build_class___obj);

py/obj.h

@@ -1015,8 +1015,6 @@ typedef struct _mp_obj_module_t {
 static inline mp_obj_dict_t *mp_obj_module_get_globals(mp_obj_t module) {
     return ((mp_obj_module_t *)MP_OBJ_TO_PTR(module))->globals;
 }
-// check if given module object is a package
-bool mp_obj_is_package(mp_obj_t module);
 
 // staticmethod and classmethod types; defined here so we can make const versions
 // this structure is used for instances of both staticmethod and classmethod

py/objmodule.c

@@ -171,35 +171,10 @@ STATIC const mp_rom_map_elem_t mp_builtin_module_table[] = {
 
 MP_DEFINE_CONST_MAP(mp_builtin_module_map, mp_builtin_module_table);
 
-// Tries to find a loaded module, otherwise attempts to load a builtin, otherwise MP_OBJ_NULL.
-mp_obj_t mp_module_get_loaded_or_builtin(qstr module_name) {
-    // First try loaded modules.
-    mp_map_elem_t *elem = mp_map_lookup(&MP_STATE_VM(mp_loaded_modules_dict).map, MP_OBJ_NEW_QSTR(module_name), MP_MAP_LOOKUP);
-
-    if (!elem) {
-        #if MICROPY_MODULE_WEAK_LINKS
-        return mp_module_get_builtin(module_name);
-        #else
-        // Otherwise try builtin.
-        elem = mp_map_lookup((mp_map_t *)&mp_builtin_module_map, MP_OBJ_NEW_QSTR(module_name), MP_MAP_LOOKUP);
-        if (!elem) {
-            return MP_OBJ_NULL;
-        }
-
-        #if MICROPY_MODULE_BUILTIN_INIT
-        // If found, it's a newly loaded built-in, so init it.
-        mp_module_call_init(MP_OBJ_NEW_QSTR(module_name), elem->value);
-        #endif
-        #endif
-    }
-
-    return elem->value;
-}
-
-#if MICROPY_MODULE_WEAK_LINKS
-// Tries to find a loaded module, otherwise attempts to load a builtin, otherwise MP_OBJ_NULL.
+// Attempts to find (and initialise) a builtin, otherwise returns MP_OBJ_NULL.
+// This must only be called after first checking the loaded modules,
+// otherwise the module will be re-initialised.
 mp_obj_t mp_module_get_builtin(qstr module_name) {
-    // Try builtin.
     mp_map_elem_t *elem = mp_map_lookup((mp_map_t *)&mp_builtin_module_map, MP_OBJ_NEW_QSTR(module_name), MP_MAP_LOOKUP);
     if (!elem) {
         return MP_OBJ_NULL;
@@ -212,7 +187,6 @@ mp_obj_t mp_module_get_builtin(qstr module_name) {
 
     return elem->value;
 }
-#endif
 
 #if MICROPY_MODULE_BUILTIN_INIT
 STATIC void mp_module_register(mp_obj_t module_name, mp_obj_t module) {

py/objmodule.h

@@ -33,10 +33,7 @@
 
 extern const mp_map_t mp_builtin_module_map;
 
-mp_obj_t mp_module_get_loaded_or_builtin(qstr module_name);
-#if MICROPY_MODULE_WEAK_LINKS
 mp_obj_t mp_module_get_builtin(qstr module_name);
-#endif
 
 void mp_module_generic_attr(qstr attr, mp_obj_t *dest, const uint16_t *keys, mp_obj_t *values);
 

py/runtime.c

@@ -1501,7 +1501,8 @@ mp_obj_t mp_import_from(mp_obj_t module, qstr name) {
     #if MICROPY_ENABLE_EXTERNAL_IMPORT
 
     // See if it's a package, then can try FS import
-    if (!mp_obj_is_package(module)) {
+    mp_load_method_maybe(module, MP_QSTR___path__, dest);
+    if (dest[0] == MP_OBJ_NULL) {
         goto import_error;
     }