py/builtinimport: Refactor module importing.

Simplify and document/comment the handling of builtin import for:
- already-loaded modules
- built-in modules
- built-in umodules (formerly weak links)
- filesystem modules

Retains existing functionality with smaller code size but should also
facilitate potential new features (built-in packages, controlling the
frozen path).

Also makes the (unix-only) -m behavior a bit more obvious and configurable.

Code size change with this commit:

   bare-arm:    +0 +0.000%
minimal x86:   -64 -0.039%
   unix x64:   -32 -0.006%
unix nanbox:    -4 -0.001%
      stm32:  -184 -0.047% PYBV10
     cc3200:  -120 -0.065%
    esp8266:  -228 -0.033% GENERIC
      esp32:  -268 -0.018% GENERIC[incl +16(data)]
        nrf:  -152 -0.087% pca10040
        rp2:  -256 -0.052% PICO
       samd:   -80 -0.057% ADAFRUIT_ITSYBITSY_M4_EXPRESS

ports/unix/mpconfigport.h

@@ -86,6 +86,7 @@
 #define MICROPY_OPT_MAP_LOOKUP_CACHE (1)
 #endif
 #define MICROPY_MODULE_WEAK_LINKS   (1)
+#define MICROPY_MODULE_OVERRIDE_MAIN_IMPORT (1)
 #define MICROPY_CAN_OVERRIDE_BUILTINS (1)
 #define MICROPY_VFS_POSIX_FILE      (1)
 #define MICROPY_PY_FUNCTION_ATTRS   (1)

ports/windows/mpconfigport.h

@@ -60,6 +60,7 @@
 #define MICROPY_STREAMS_POSIX_API   (1)
 #define MICROPY_OPT_COMPUTED_GOTO   (0)
 #define MICROPY_MODULE_WEAK_LINKS   (1)
+#define MICROPY_MODULE_OVERRIDE_MAIN_IMPORT (1)
 #define MICROPY_CAN_OVERRIDE_BUILTINS (1)
 #define MICROPY_VFS_POSIX_FILE      (1)
 #define MICROPY_PY_FUNCTION_ATTRS   (1)

py/builtinimport.c

@@ -5,6 +5,7 @@
  *
  * Copyright (c) 2013-2019 Damien P. George
  * Copyright (c) 2014 Paul Sokolovsky
+ * Copyright (c) 2021 Jim Mussared
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -46,7 +47,8 @@
 
 #if MICROPY_ENABLE_EXTERNAL_IMPORT
 
-#define PATH_SEP_CHAR '/'
+// Must be a string of one byte.
+#define PATH_SEP_CHAR "/"
 
 bool mp_obj_is_package(mp_obj_t module) {
     mp_obj_t dest[2];
@@ -54,9 +56,11 @@ bool mp_obj_is_package(mp_obj_t module) {
     return dest[0] != MP_OBJ_NULL;
 }
 
-// Stat either frozen or normal module by a given path
-// (whatever is available, if at all).
-STATIC mp_import_stat_t mp_import_stat_any(const char *path) {
+// Wrapper for mp_import_stat (which is provided by the port, and typically
+// uses mp_vfs_import_stat) to also search frozen modules. Given an exact
+// path to a file or directory (e.g. "foo/bar", foo/bar.py" or "foo/bar.mpy"),
+// will return whether the path is a file, directory, or doesn't exist.
+STATIC mp_import_stat_t stat_path_or_frozen(const char *path) {
     #if MICROPY_MODULE_FROZEN
     mp_import_stat_t st = mp_frozen_stat(path);
     if (st != MP_IMPORT_STAT_NO_EXIST) {
@@ -66,15 +70,18 @@ STATIC mp_import_stat_t mp_import_stat_any(const char *path) {
     return mp_import_stat(path);
 }
 
+// Given a path to a .py file, try and find this path as either a .py or .mpy
+// in either the filesystem or frozen modules.
 STATIC mp_import_stat_t stat_file_py_or_mpy(vstr_t *path) {
-    mp_import_stat_t stat = mp_import_stat_any(vstr_null_terminated_str(path));
+    mp_import_stat_t stat = stat_path_or_frozen(vstr_null_terminated_str(path));
     if (stat == MP_IMPORT_STAT_FILE) {
         return stat;
     }
 
     #if MICROPY_PERSISTENT_CODE_LOAD
+    // Didn't find .py -- try the .mpy instead by inserting an 'm' into the '.py'.
     vstr_ins_byte(path, path->len - 2, 'm');
-    stat = mp_import_stat_any(vstr_null_terminated_str(path));
+    stat = stat_path_or_frozen(vstr_null_terminated_str(path));
     if (stat == MP_IMPORT_STAT_FILE) {
         return stat;
     }
@@ -83,8 +90,10 @@ STATIC mp_import_stat_t stat_file_py_or_mpy(vstr_t *path) {
     return MP_IMPORT_STAT_NO_EXIST;
 }
 
+// Given an import path (e.g. "foo/bar"), try and find "foo/bar" (a directory)
+// or "foo/bar.(m)py" in either the filesystem or frozen modules.
 STATIC mp_import_stat_t stat_dir_or_file(vstr_t *path) {
-    mp_import_stat_t stat = mp_import_stat_any(vstr_null_terminated_str(path));
+    mp_import_stat_t stat = stat_path_or_frozen(vstr_null_terminated_str(path));
     DEBUG_printf("stat %s: %d\n", vstr_str(path), stat);
     if (stat == MP_IMPORT_STAT_DIR) {
         return stat;
@@ -95,14 +104,16 @@ STATIC mp_import_stat_t stat_dir_or_file(vstr_t *path) {
     return stat_file_py_or_mpy(path);
 }
 
-STATIC mp_import_stat_t find_file(const char *file_str, uint file_len, vstr_t *dest) {
+// Given a top-level module, try and find it in each of the sys.path entries
+// via stat_dir_or_file.
+STATIC mp_import_stat_t stat_top_level_dir_or_file(qstr mod_name, vstr_t *dest) {
+    DEBUG_printf("stat_top_level_dir_or_file: '%s'\n", qstr_str(mod_name));
     #if MICROPY_PY_SYS
-    // extract the list of paths
     size_t path_num;
     mp_obj_t *path_items;
     mp_obj_list_get(mp_sys_path, &path_num, &path_items);
 
-    if (path_num != 0) {
+    if (path_num > 0) {
         // go through each path looking for a directory or file
         for (size_t i = 0; i < path_num; i++) {
             vstr_reset(dest);
@@ -110,9 +121,9 @@ STATIC mp_import_stat_t find_file(const char *file_str, uint file_len, vstr_t *d
             const char *p = mp_obj_str_get_data(path_items[i], &p_len);
             if (p_len > 0) {
                 vstr_add_strn(dest, p, p_len);
-                vstr_add_char(dest, PATH_SEP_CHAR);
+                vstr_add_char(dest, PATH_SEP_CHAR[0]);
             }
-            vstr_add_strn(dest, file_str, file_len);
+            vstr_add_str(dest, qstr_str(mod_name));
             mp_import_stat_t stat = stat_dir_or_file(dest);
             if (stat != MP_IMPORT_STAT_NO_EXIST) {
                 return stat;
@@ -124,8 +135,9 @@ STATIC mp_import_stat_t find_file(const char *file_str, uint file_len, vstr_t *d
     }
     #endif
 
-    // mp_sys_path is empty, so just use the given file name
-    vstr_add_strn(dest, file_str, file_len);
+    // mp_sys_path is empty (or not enabled), so just stat the given path
+    // directly.
+    vstr_add_str(dest, qstr_str(mod_name));
     return stat_dir_or_file(dest);
 }
 
@@ -189,7 +201,6 @@ STATIC void do_load(mp_obj_t module_obj, vstr_t *file) {
     #if MICROPY_MODULE_FROZEN
     void *modref;
     int frozen_type = mp_find_frozen_module(file_str, file->len, &modref);
-    #endif
 
     // If we support frozen str modules and the compiler is enabled, and we
     // found the filename in the list of frozen files, then load and execute it.
@@ -209,6 +220,8 @@ STATIC void do_load(mp_obj_t module_obj, vstr_t *file) {
     }
     #endif
 
+    #endif // MICROPY_MODULE_FROZEN
+
     // If we support loading .mpy files then check if the file extension is of
     // the correct format and, if so, load and execute the file.
     #if MICROPY_HAS_FILE_READER && MICROPY_PERSISTENT_CODE_LOAD
@@ -232,15 +245,212 @@ STATIC void do_load(mp_obj_t module_obj, vstr_t *file) {
     #endif
 }
 
-STATIC void chop_component(const char *start, const char **end) {
-    const char *p = *end;
-    while (p > start) {
+// Convert a relative (to the current module) import, going up "level" levels,
+// into an absolute import.
+STATIC void evaluate_relative_import(mp_int_t level, const char **module_name, size_t *module_name_len) {
+    // What we want to do here is to take the name of the current module,
+    // remove <level> trailing components, and concatenate the passed-in
+    // module name.
+    // For example, level=3, module_name="foo.bar", __name__="a.b.c.d" --> "a.foo.bar"
+    // "Relative imports use a module's __name__ attribute to determine that
+    // module's position in the package hierarchy."
+    // http://legacy.python.org/dev/peps/pep-0328/#relative-imports-and-name
+
+    mp_obj_t current_module_name_obj = mp_obj_dict_get(MP_OBJ_FROM_PTR(mp_globals_get()), MP_OBJ_NEW_QSTR(MP_QSTR___name__));
+    assert(current_module_name_obj != MP_OBJ_NULL);
+
+    #if MICROPY_MODULE_OVERRIDE_MAIN_IMPORT && MICROPY_CPYTHON_COMPAT
+    if (MP_OBJ_QSTR_VALUE(current_module_name_obj) == MP_QSTR___main__) {
+        // This is a module loaded by -m command-line switch (e.g. unix port),
+        // and so its __name__ has been set to "__main__". Get its real name
+        // that we stored during import in the __main__ attribute.
+        current_module_name_obj = mp_obj_dict_get(MP_OBJ_FROM_PTR(mp_globals_get()), MP_OBJ_NEW_QSTR(MP_QSTR___main__));
+    }
+    #endif
+
+    // If we have a __path__ in the globals dict, then we're a package.
+    bool is_pkg = mp_map_lookup(&mp_globals_get()->map, MP_OBJ_NEW_QSTR(MP_QSTR___path__), MP_MAP_LOOKUP);
+
+    #if DEBUG_PRINT
+    DEBUG_printf("Current module/package: ");
+    mp_obj_print_helper(MICROPY_DEBUG_PRINTER, current_module_name_obj, PRINT_REPR);
+    DEBUG_printf(", is_package: %d", is_pkg);
+    DEBUG_printf("\n");
+    #endif
+
+    size_t current_module_name_len;
+    const char *current_module_name = mp_obj_str_get_data(current_module_name_obj, &current_module_name_len);
+
+    const char *p = current_module_name + current_module_name_len;
+    if (is_pkg) {
+        // If we're evaluating relative to a package, then take off one fewer
+        // level (i.e. the relative search starts inside the package, rather
+        // than as a sibling of the package).
+        --level;
+    }
+
+    // Walk back 'level' dots (or run out of path).
+    while (level && p > current_module_name) {
         if (*--p == '.') {
-            *end = p;
-            return;
+            --level;
         }
     }
-    *end = p;
+
+    // We must have some component left over to import from.
+    if (p == current_module_name) {
+        mp_raise_msg(&mp_type_ImportError, MP_ERROR_TEXT("can't perform relative import"));
+    }
+
+    // New length is len("<chopped path>.<module_name>"). Note: might be one byte
+    // more than we need if module_name is empty (for the extra . we will
+    // append).
+    uint new_module_name_len = (size_t)(p - current_module_name) + 1 + *module_name_len;
+    char *new_mod = mp_local_alloc(new_module_name_len);
+    memcpy(new_mod, current_module_name, p - current_module_name);
+
+    // Only append ".<module_name>" if there was one).
+    if (*module_name_len != 0) {
+        new_mod[p - current_module_name] = '.';
+        memcpy(new_mod + (p - current_module_name) + 1, *module_name, *module_name_len);
+    } else {
+        --new_module_name_len;
+    }
+
+    // Copy into a QSTR.
+    qstr new_mod_q = qstr_from_strn(new_mod, new_module_name_len);
+    mp_local_free(new_mod);
+
+    DEBUG_printf("Resolved base name for relative import: '%s'\n", qstr_str(new_mod_q));
+    *module_name = qstr_str(new_mod_q);
+    *module_name_len = new_module_name_len;
+}
+
+// Load a module at the specified absolute path, possibly as a submodule of the given outer module.
+// full_mod_name:    The full absolute path to this module (e.g. "foo.bar.baz").
+// level_mod_name:   The final component of the path (e.g. "baz").
+// outer_module_obj: The parent module (we need to store this module as an
+//                   attribute on it) (or MP_OBJ_NULL for top-level).
+// path:             The filesystem path where we found the parent module
+//                   (or empty for a top level module).
+// override_main:    Whether to set the __name__ to "__main__" (and use __main__
+//                   for the actual path).
+STATIC mp_obj_t process_import_at_level(qstr full_mod_name, qstr level_mod_name, mp_obj_t outer_module_obj, vstr_t *path, bool override_main) {
+    mp_import_stat_t stat = MP_IMPORT_STAT_NO_EXIST;
+
+    // Exact-match of built-in (or already-loaded) takes priority.
+    mp_obj_t module_obj = mp_module_get_loaded_or_builtin(full_mod_name);
+
+    // Even if we find the module, go through the motions of searching for it
+    // because we may actually be in the process of importing a sub-module.
+    // So we need to (re-)find the correct path to be finding the sub-module
+    // on the next iteration of process_import_at_level.
+
+    if (outer_module_obj == MP_OBJ_NULL) {
+        DEBUG_printf("Searching for top-level module\n");
+
+        // First module in the dotted-name; search for a directory or file
+        // relative to all the locations in sys.path.
+        stat = stat_top_level_dir_or_file(full_mod_name, path);
+
+        // If the module "foo" doesn't exist on the filesystem, and it's not a
+        // builtin, try and find "ufoo" as a built-in. (This feature was
+        // formerly known as "weak links").
+        #if MICROPY_MODULE_WEAK_LINKS
+        if (stat == MP_IMPORT_STAT_NO_EXIST && module_obj == MP_OBJ_NULL) {
+            char *umodule_buf = vstr_str(path);
+            umodule_buf[0] = 'u';
+            strcpy(umodule_buf + 1, qstr_str(level_mod_name));
+            qstr umodule_name = qstr_from_str(umodule_buf);
+            module_obj = mp_module_get_builtin(umodule_name);
+        }
+        #endif
+    } else {
+        DEBUG_printf("Searching for sub-module\n");
+
+        // Add the current part of the module name to the path.
+        vstr_add_char(path, PATH_SEP_CHAR[0]);
+        vstr_add_str(path, qstr_str(level_mod_name));
+
+        // Because it's not top level, we already know which path the parent was found in.
+        stat = stat_dir_or_file(path);
+    }
+    DEBUG_printf("Current path: %.*s\n", (int)vstr_len(path), vstr_str(path));
+
+    if (module_obj == MP_OBJ_NULL) {
+        // Not a built-in and not already-loaded.
+
+        if (stat == MP_IMPORT_STAT_NO_EXIST) {
+            // And the file wasn't found -- fail.
+            #if MICROPY_ERROR_REPORTING <= MICROPY_ERROR_REPORTING_TERSE
+            mp_raise_msg(&mp_type_ImportError, MP_ERROR_TEXT("module not found"));
+            #else
+            mp_raise_msg_varg(&mp_type_ImportError, MP_ERROR_TEXT("no module named '%q'"), full_mod_name);
+            #endif
+        }
+
+        // Not a built-in but found on the filesystem, try and load it.
+
+        DEBUG_printf("Found path: %.*s\n", (int)vstr_len(path), vstr_str(path));
+
+        // Prepare for loading from the filesystem. Create a new shell module.
+        module_obj = mp_obj_new_module(full_mod_name);
+
+        #if MICROPY_MODULE_OVERRIDE_MAIN_IMPORT
+        // If this module is being loaded via -m on unix, then
+        // override __name__ to "__main__". Do this only for *modules*
+        // however - packages never have their names replaced, instead
+        // they're -m'ed using a special __main__ submodule in them. (This all
+        // apparently is done to not touch the package name itself, which is
+        // important for future imports).
+        if (override_main && stat != MP_IMPORT_STAT_DIR) {
+            mp_obj_module_t *o = MP_OBJ_TO_PTR(module_obj);
+            mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR___main__));
+            #if MICROPY_CPYTHON_COMPAT
+            // Store module as "__main__" in the dictionary of loaded modules (returned by sys.modules).
+            mp_obj_dict_store(MP_OBJ_FROM_PTR(&MP_STATE_VM(mp_loaded_modules_dict)), MP_OBJ_NEW_QSTR(MP_QSTR___main__), module_obj);
+            // Store real name in "__main__" attribute. Need this for
+            // resolving relative imports later. "__main__ was chosen
+            // semi-randonly, to reuse existing qstr's.
+            mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___main__), MP_OBJ_NEW_QSTR(full_mod_name));
+            #endif
+        }
+        #endif // MICROPY_MODULE_OVERRIDE_MAIN_IMPORT
+
+        if (stat == MP_IMPORT_STAT_DIR) {
+            // Directory -- execute "path/__init__.py".
+            DEBUG_printf("%.*s is dir\n", (int)vstr_len(path), vstr_str(path));
+            // Store the __path__ attribute onto this module.
+            // https://docs.python.org/3/reference/import.html
+            // "Specifically, any module that contains a __path__ attribute is considered a package."
+            mp_store_attr(module_obj, MP_QSTR___path__, mp_obj_new_str(vstr_str(path), vstr_len(path)));
+            size_t orig_path_len = path->len;
+            vstr_add_str(path, PATH_SEP_CHAR "__init__.py");
+            if (stat_file_py_or_mpy(path) == MP_IMPORT_STAT_FILE) {
+                do_load(module_obj, path);
+            } else {
+                // No-op. Nothing to load.
+                // mp_warning("%s is imported as namespace package", vstr_str(&path));
+            }
+            // Remove /__init__.py suffix.
+            path->len = orig_path_len;
+        } else { // MP_IMPORT_STAT_FILE
+            // File -- execute "path.(m)py".
+            do_load(module_obj, path);
+            // Note: This should be the last component in the import path.  If
+            // there are remaining components then it's an ImportError
+            // because the current path(the module that was just loaded) is
+            // not a package.  This will be caught on the next iteration
+            // because the file will not exist.
+        }
+    }
+
+    if (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.
+        mp_store_attr(outer_module_obj, level_mod_name, module_obj);
+    }
+
+    return module_obj;
 }
 
 mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
@@ -253,9 +463,23 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
     }
     #endif
 
-    mp_obj_t module_name = args[0];
+    // This is the import path, with any leading dots stripped.
+    // "import foo.bar" --> module_name="foo.bar"
+    // "from foo.bar import baz" --> module_name="foo.bar"
+    // "from . import foo" --> module_name=""
+    // "from ...foo.bar import baz" --> module_name="foo.bar"
+    mp_obj_t module_name_obj = args[0];
+
+    // These are the imported names.
+    // i.e. "from foo.bar import baz, zap" --> fromtuple=("baz", "zap",)
+    // Note: There's a special case on the Unix port, where this is set to mp_const_false which means that it's __main__.
     mp_obj_t fromtuple = mp_const_none;
+
+    // Level is the number of leading dots in a relative import.
+    // i.e. "from . import foo" --> level=1
+    // i.e. "from ...foo.bar import baz" --> level=3
     mp_int_t level = 0;
+
     if (n_args >= 4) {
         fromtuple = args[3];
         if (n_args >= 5) {
@@ -266,211 +490,64 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
         }
     }
 
-    size_t mod_len;
-    const char *mod_str = mp_obj_str_get_data(module_name, &mod_len);
+    size_t module_name_len;
+    const char *module_name = mp_obj_str_get_data(module_name_obj, &module_name_len);
 
     if (level != 0) {
-        // What we want to do here is to take name of current module,
-        // chop <level> trailing components, and concatenate with passed-in
-        // module name, thus resolving relative import name into absolute.
-        // This even appears to be correct per
-        // http://legacy.python.org/dev/peps/pep-0328/#relative-imports-and-name
-        // "Relative imports use a module's __name__ attribute to determine that
-        // module's position in the package hierarchy."
-        level--;
-        mp_obj_t this_name_q = mp_obj_dict_get(MP_OBJ_FROM_PTR(mp_globals_get()), MP_OBJ_NEW_QSTR(MP_QSTR___name__));
-        assert(this_name_q != MP_OBJ_NULL);
-        #if MICROPY_CPYTHON_COMPAT
-        if (MP_OBJ_QSTR_VALUE(this_name_q) == MP_QSTR___main__) {
-            // This is a module run by -m command-line switch, get its real name from backup attribute
-            this_name_q = mp_obj_dict_get(MP_OBJ_FROM_PTR(mp_globals_get()), MP_OBJ_NEW_QSTR(MP_QSTR___main__));
-        }
-        #endif
-        mp_map_t *globals_map = &mp_globals_get()->map;
-        mp_map_elem_t *elem = mp_map_lookup(globals_map, MP_OBJ_NEW_QSTR(MP_QSTR___path__), MP_MAP_LOOKUP);
-        bool is_pkg = (elem != NULL);
-
-        #if DEBUG_PRINT
-        DEBUG_printf("Current module/package: ");
-        mp_obj_print_helper(MICROPY_DEBUG_PRINTER, this_name_q, PRINT_REPR);
-        DEBUG_printf(", is_package: %d", is_pkg);
-        DEBUG_printf("\n");
-        #endif
-
-        size_t this_name_l;
-        const char *this_name = mp_obj_str_get_data(this_name_q, &this_name_l);
-
-        const char *p = this_name + this_name_l;
-        if (!is_pkg) {
-            // We have module, but relative imports are anchored at package, so
-            // go there.
-            chop_component(this_name, &p);
-        }
-
-        while (level--) {
-            chop_component(this_name, &p);
-        }
-
-        // We must have some component left over to import from
-        if (p == this_name) {
-            mp_raise_msg(&mp_type_ImportError, MP_ERROR_TEXT("can't perform relative import"));
-        }
-
-        uint new_mod_l = (mod_len == 0 ? (size_t)(p - this_name) : (size_t)(p - this_name) + 1 + mod_len);
-        char *new_mod = mp_local_alloc(new_mod_l);
-        memcpy(new_mod, this_name, p - this_name);
-        if (mod_len != 0) {
-            new_mod[p - this_name] = '.';
-            memcpy(new_mod + (p - this_name) + 1, mod_str, mod_len);
-        }
-
-        qstr new_mod_q = qstr_from_strn(new_mod, new_mod_l);
-        mp_local_free(new_mod);
-        DEBUG_printf("Resolved base name for relative import: '%s'\n", qstr_str(new_mod_q));
-        module_name = MP_OBJ_NEW_QSTR(new_mod_q);
-        mod_str = qstr_str(new_mod_q);
-        mod_len = new_mod_l;
+        // Turn "foo.bar" into "<current module minus 3 components>.foo.bar".
+        evaluate_relative_import(level, &module_name, &module_name_len);
     }
 
-    if (mod_len == 0) {
+    if (module_name_len == 0) {
         mp_raise_ValueError(NULL);
     }
 
-    // check if module already exists
-    qstr module_name_qstr = mp_obj_str_get_qstr(module_name);
-    mp_obj_t module_obj = mp_module_get(module_name_qstr);
-    if (module_obj != MP_OBJ_NULL) {
-        DEBUG_printf("Module already loaded\n");
-        // If it's not a package, return module right away
-        char *p = strchr(mod_str, '.');
-        if (p == NULL) {
-            return module_obj;
-        }
-        // If fromlist is not empty, return leaf module
-        if (fromtuple != mp_const_none) {
-            return module_obj;
-        }
-        // Otherwise, we need to return top-level package
-        qstr pkg_name = qstr_from_strn(mod_str, p - mod_str);
-        return mp_module_get(pkg_name);
-    }
-    DEBUG_printf("Module not yet loaded\n");
+    DEBUG_printf("Starting module search for '%s'\n", module_name);
 
-    uint last = 0;
     VSTR_FIXED(path, MICROPY_ALLOC_PATH_MAX)
-    module_obj = MP_OBJ_NULL;
     mp_obj_t top_module_obj = MP_OBJ_NULL;
     mp_obj_t outer_module_obj = MP_OBJ_NULL;
-    uint i;
-    for (i = 1; i <= mod_len; i++) {
-        if (i == mod_len || mod_str[i] == '.') {
-            // create a qstr for the module name up to this depth
-            qstr mod_name = qstr_from_strn(mod_str, i);
-            DEBUG_printf("Processing module: %s\n", qstr_str(mod_name));
-            DEBUG_printf("Previous path: =%.*s=\n", vstr_len(&path), vstr_str(&path));
 
-            // find the file corresponding to the module name
-            mp_import_stat_t stat;
-            if (vstr_len(&path) == 0) {
-                // first module in the dotted-name; search for a directory or file
-                stat = find_file(mod_str, i, &path);
-            } else {
-                // latter module in the dotted-name; append to path
-                vstr_add_char(&path, PATH_SEP_CHAR);
-                vstr_add_strn(&path, mod_str + last, i - last);
-                stat = stat_dir_or_file(&path);
-            }
-            DEBUG_printf("Current path: %.*s\n", vstr_len(&path), vstr_str(&path));
+    // Search for the end of each component.
+    size_t current_component_start = 0;
+    for (size_t i = 1; i <= module_name_len; i++) {
+        if (i == module_name_len || module_name[i] == '.') {
+            // The module name up to this depth (e.g. foo.bar.baz).
+            qstr full_mod_name = qstr_from_strn(module_name, i);
+            // The current level name (e.g. baz).
+            qstr level_mod_name = qstr_from_strn(module_name + current_component_start, i - current_component_start);
 
-            if (stat == MP_IMPORT_STAT_NO_EXIST) {
-                module_obj = MP_OBJ_NULL;
-                #if MICROPY_MODULE_WEAK_LINKS
-                // check if there is a weak link to this module
-                if (i == mod_len) {
-                    module_obj = mp_module_search_umodule(mod_str);
-                    if (module_obj != MP_OBJ_NULL) {
-                        // found weak linked module
-                        mp_module_call_init(mod_name, module_obj);
-                    }
-                }
-                #endif
-                if (module_obj == MP_OBJ_NULL) {
-                    // couldn't find the file, so fail
-                    #if MICROPY_ERROR_REPORTING <= MICROPY_ERROR_REPORTING_TERSE
-                    mp_raise_msg(&mp_type_ImportError, MP_ERROR_TEXT("module not found"));
-                    #else
-                    mp_raise_msg_varg(&mp_type_ImportError, MP_ERROR_TEXT("no module named '%q'"), mod_name);
-                    #endif
-                }
-            } else {
-                // found the file, so get the module
-                module_obj = mp_module_get(mod_name);
-            }
+            DEBUG_printf("Processing module: '%s' at level '%s'\n", qstr_str(full_mod_name), qstr_str(level_mod_name));
+            DEBUG_printf("Previous path: =%.*s=\n", (int)vstr_len(&path), vstr_str(&path));
 
-            if (module_obj == MP_OBJ_NULL) {
-                // module not already loaded, so load it!
+            #if MICROPY_MODULE_OVERRIDE_MAIN_IMPORT
+            // On unix, if this is being loaded via -m (magic mp_const_false),
+            // then handle that if it's the final component.
+            bool override_main = (i == module_name_len && fromtuple == mp_const_false);
+            #else
+            bool override_main = false;
+            #endif
 
-                module_obj = mp_obj_new_module(mod_name);
+            // Import this module.
+            mp_obj_t module_obj = process_import_at_level(full_mod_name, level_mod_name, outer_module_obj, &path, override_main);
 
-                // if args[3] (fromtuple) has magic value False, set up
-                // this module for command-line "-m" option (set module's
-                // name to __main__ instead of real name). Do this only
-                // for *modules* however - packages never have their names
-                // replaced, instead they're -m'ed using a special __main__
-                // submodule in them. (This all apparently is done to not
-                // touch package name itself, which is important for future
-                // imports).
-                if (i == mod_len && fromtuple == mp_const_false && stat != MP_IMPORT_STAT_DIR) {
-                    mp_obj_module_t *o = MP_OBJ_TO_PTR(module_obj);
-                    mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___name__), MP_OBJ_NEW_QSTR(MP_QSTR___main__));
-                    #if MICROPY_CPYTHON_COMPAT
-                    // Store module as "__main__" in the dictionary of loaded modules (returned by sys.modules).
-                    mp_obj_dict_store(MP_OBJ_FROM_PTR(&MP_STATE_VM(mp_loaded_modules_dict)), MP_OBJ_NEW_QSTR(MP_QSTR___main__), module_obj);
-                    // Store real name in "__main__" attribute. Chosen semi-randonly, to reuse existing qstr's.
-                    mp_obj_dict_store(MP_OBJ_FROM_PTR(o->globals), MP_OBJ_NEW_QSTR(MP_QSTR___main__), MP_OBJ_NEW_QSTR(mod_name));
-                    #endif
-                }
-
-                if (stat == MP_IMPORT_STAT_DIR) {
-                    DEBUG_printf("%.*s is dir\n", vstr_len(&path), vstr_str(&path));
-                    // https://docs.python.org/3/reference/import.html
-                    // "Specifically, any module that contains a __path__ attribute is considered a package."
-                    mp_store_attr(module_obj, MP_QSTR___path__, mp_obj_new_str(vstr_str(&path), vstr_len(&path)));
-                    size_t orig_path_len = path.len;
-                    vstr_add_char(&path, PATH_SEP_CHAR);
-                    vstr_add_str(&path, "__init__.py");
-                    if (stat_file_py_or_mpy(&path) != MP_IMPORT_STAT_FILE) {
-                        // mp_warning("%s is imported as namespace package", vstr_str(&path));
-                    } else {
-                        do_load(module_obj, &path);
-                    }
-                    path.len = orig_path_len;
-                } else { // MP_IMPORT_STAT_FILE
-                    do_load(module_obj, &path);
-                    // This should be the last component in the import path.  If there are
-                    // remaining components then it's an ImportError because the current path
-                    // (the module that was just loaded) is not a package.  This will be caught
-                    // on the next iteration because the file will not exist.
-                }
-            }
-            if (outer_module_obj != MP_OBJ_NULL) {
-                qstr s = qstr_from_strn(mod_str + last, i - last);
-                mp_store_attr(outer_module_obj, s, module_obj);
-            }
+            // Set this as the parent module, and remember the top-level module if it's the first.
             outer_module_obj = module_obj;
             if (top_module_obj == MP_OBJ_NULL) {
                 top_module_obj = module_obj;
             }
-            last = i + 1;
+
+            current_component_start = i + 1;
         }
     }
 
-    // If fromlist is not empty, return leaf module
     if (fromtuple != mp_const_none) {
-        return module_obj;
+        // If fromtuple is not empty, return leaf module
+        return outer_module_obj;
+    } else {
+        // Otherwise, we need to return top-level package
+        return top_module_obj;
     }
-    // Otherwise, we need to return top-level package
-    return top_module_obj;
 }
 
 #else // MICROPY_ENABLE_EXTERNAL_IMPORT
@@ -483,17 +560,19 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
 
     // Check if module already exists, and return it if it does
     qstr module_name_qstr = mp_obj_str_get_qstr(args[0]);
-    mp_obj_t module_obj = mp_module_get(module_name_qstr);
+    mp_obj_t module_obj = mp_module_get_loaded_or_builtin(module_name_qstr);
     if (module_obj != MP_OBJ_NULL) {
         return module_obj;
     }
 
     #if MICROPY_MODULE_WEAK_LINKS
     // Check if there is a weak link to this module
-    module_obj = mp_module_search_umodule(qstr_str(module_name_qstr));
+    char umodule_buf[MICROPY_ALLOC_PATH_MAX];
+    umodule_buf[0] = 'u';
+    strcpy(umodule_buf + 1, args[0]);
+    qstr umodule_name_qstr = qstr_from_str(umodule_buf);
+    module_obj = mp_module_get_loaded_or_builtin(umodule_name_qstr);
     if (module_obj != MP_OBJ_NULL) {
-        // Found weak-linked module
-        mp_module_call_init(module_name_qstr, module_obj);
         return module_obj;
     }
     #endif

py/mpconfig.h

@@ -829,6 +829,12 @@ typedef double mp_float_t;
 #define MICROPY_MODULE_WEAK_LINKS (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
 #endif
 
+// Whether to enable importing foo.py with __name__ set to '__main__'
+// Used by the unix port for the -m flag.
+#ifndef MICROPY_MODULE_OVERRIDE_MAIN_IMPORT
+#define MICROPY_MODULE_OVERRIDE_MAIN_IMPORT (0)
+#endif
+
 // Whether frozen modules are supported in the form of strings
 #ifndef MICROPY_MODULE_FROZEN_STR
 #define MICROPY_MODULE_FROZEN_STR (0)

py/objmodule.c

@@ -35,6 +35,10 @@
 
 #include "genhdr/moduledefs.h"
 
+#if MICROPY_MODULE_BUILTIN_INIT
+STATIC void mp_module_call_init(mp_obj_t module_name, mp_obj_t module_obj);
+#endif
+
 STATIC void module_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
     (void)kind;
     mp_obj_module_t *self = MP_OBJ_TO_PTR(self_in);
@@ -245,47 +249,56 @@ STATIC const mp_rom_map_elem_t mp_builtin_module_table[] = {
 
 MP_DEFINE_CONST_MAP(mp_builtin_module_map, mp_builtin_module_table);
 
-// returns MP_OBJ_NULL if not found
-mp_obj_t mp_module_get(qstr module_name) {
-    mp_map_t *mp_loaded_modules_map = &MP_STATE_VM(mp_loaded_modules_dict).map;
-    // lookup module
-    mp_map_elem_t *el = mp_map_lookup(mp_loaded_modules_map, MP_OBJ_NEW_QSTR(module_name), MP_MAP_LOOKUP);
+// 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 (el == NULL) {
-        // module not found, look for builtin module names
-        el = mp_map_lookup((mp_map_t *)&mp_builtin_module_map, MP_OBJ_NEW_QSTR(module_name), MP_MAP_LOOKUP);
-        if (el == NULL) {
+    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;
         }
-        mp_module_call_init(module_name, el->value);
+
+        #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
     }
 
-    // module found, return it
-    return el->value;
-}
-
-void mp_module_register(qstr qst, mp_obj_t module) {
-    mp_map_t *mp_loaded_modules_map = &MP_STATE_VM(mp_loaded_modules_dict).map;
-    mp_map_lookup(mp_loaded_modules_map, MP_OBJ_NEW_QSTR(qst), MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = module;
+    return elem->value;
 }
 
 #if MICROPY_MODULE_WEAK_LINKS
-// Search for u"foo" in built-in modules, return MP_OBJ_NULL if not found
-mp_obj_t mp_module_search_umodule(const char *module_str) {
-    for (size_t i = 0; i < MP_ARRAY_SIZE(mp_builtin_module_table); ++i) {
-        const mp_map_elem_t *entry = (const mp_map_elem_t *)&mp_builtin_module_table[i];
-        const char *key = qstr_str(MP_OBJ_QSTR_VALUE(entry->key));
-        if (key[0] == 'u' && strcmp(&key[1], module_str) == 0) {
-            return (mp_obj_t)entry->value;
-        }
-
+// Tries to find a loaded module, otherwise attempts to load a builtin, otherwise MP_OBJ_NULL.
+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;
     }
-    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
+
+    return elem->value;
 }
 #endif
 
 #if MICROPY_MODULE_BUILTIN_INIT
-void mp_module_call_init(qstr module_name, mp_obj_t module_obj) {
+STATIC void mp_module_register(mp_obj_t module_name, mp_obj_t module) {
+    mp_map_t *mp_loaded_modules_map = &MP_STATE_VM(mp_loaded_modules_dict).map;
+    mp_map_lookup(mp_loaded_modules_map, module_name, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = module;
+}
+
+STATIC void mp_module_call_init(mp_obj_t module_name, mp_obj_t module_obj) {
     // Look for __init__ and call it if it exists
     mp_obj_t dest[2];
     mp_load_method_maybe(module_obj, MP_QSTR___init__, dest);

py/objmodule.h

@@ -30,18 +30,9 @@
 
 extern const mp_map_t mp_builtin_module_map;
 
-mp_obj_t mp_module_get(qstr module_name);
-void mp_module_register(qstr qstr, mp_obj_t module);
-
-mp_obj_t mp_module_search_umodule(const char *module_str);
-
-#if MICROPY_MODULE_BUILTIN_INIT
-void mp_module_call_init(qstr module_name, mp_obj_t module_obj);
-#else
-static inline void mp_module_call_init(qstr module_name, mp_obj_t module_obj) {
-    (void)module_name;
-    (void)module_obj;
-}
+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
 
 #endif // MICROPY_INCLUDED_PY_OBJMODULE_H