py/makemoduledefs.py: Add a way to register extensible built-in modules.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>

py/builtinhelp.c

@@ -79,6 +79,7 @@ STATIC void mp_help_print_modules(void) {
     mp_obj_t list = mp_obj_new_list(0, NULL);
 
     mp_help_add_from_map(list, &mp_builtin_module_map);
+    mp_help_add_from_map(list, &mp_builtin_extensible_module_map);
 
     #if MICROPY_MODULE_FROZEN
     extern const char mp_frozen_names[];

py/builtinimport.c

@@ -384,28 +384,22 @@ STATIC mp_obj_t process_import_at_level(qstr full_mod_name, qstr level_mod_name,
 
         // An import of a non-extensible built-in will always bypass the
         // filesystem. e.g. `import micropython` or `import pyb`.
-        module_obj = mp_module_get_builtin(level_mod_name);
+        module_obj = mp_module_get_builtin(level_mod_name, false);
         if (module_obj != MP_OBJ_NULL) {
             return module_obj;
         }
 
-        #if MICROPY_PY_SYS
-        // Never allow sys to be overridden from the filesystem. If weak links
-        // are disabled, then this also provides a default weak link so that
-        // `import sys` is treated like `import usys` (and therefore bypasses
-        // the filesystem).
-        if (level_mod_name == MP_QSTR_sys) {
-            return MP_OBJ_FROM_PTR(&mp_module_sys);
-        }
-        #endif
-
         // First module in the dotted-name; search for a directory or file
         // relative to all the locations in sys.path.
         stat = stat_top_level(level_mod_name, &path);
 
         // TODO: If stat failed, now try extensible built-in modules.
-
+        if (stat == MP_IMPORT_STAT_NO_EXIST) {
-        // TODO: If importing `ufoo`, try `foo`.
+            module_obj = mp_module_get_builtin(level_mod_name, true);
+            if (module_obj != MP_OBJ_NULL) {
+                return module_obj;
+            }
+        }
     } else {
         DEBUG_printf("Searching for sub-module\n");
 
@@ -637,7 +631,7 @@ mp_obj_t mp_builtin___import___default(size_t n_args, const mp_obj_t *args) {
 
     // Try the name directly as a built-in.
     qstr module_name_qstr = mp_obj_str_get_qstr(args[0]);
-    mp_obj_t module_obj = mp_module_get_builtin(module_name_qstr);
+    mp_obj_t module_obj = mp_module_get_builtin(module_name_qstr, false);
     if (module_obj != MP_OBJ_NULL) {
         return module_obj;
     }

py/makemoduledefs.py

@@ -1,8 +1,17 @@
 """
 This pre-processor parses a single file containing a list of
-MP_REGISTER_MODULE(module_name, obj_module)
+`MP_REGISTER_MODULE(MP_QSTR_module_name, obj_module)` or
-These are used to generate a header with the required entries for
+`MP_REGISTER_EXTENSIBLE_MODULE(MP_QSTR_module_name, obj_module)`
-"mp_rom_map_elem_t mp_builtin_module_table[]" in py/objmodule.c
+(i.e. the output of `py/makeqstrdefs.py cat module`).
+
+The output is a header (typically moduledefs.h) which is included by
+py/objmodule.c that contains entries to be included in the definition of
+ - mp_rom_map_elem_t mp_builtin_module_table[]
+ - mp_rom_map_elem_t mp_builtin_extensible_module_table[]
+
+Extensible modules are modules that can be overridden from the filesystem, see
+py/builtinimnport.c:process_import_at_level. Regular modules will always use
+the built-in version.
 """
 
 from __future__ import print_function
@@ -13,7 +22,10 @@ import io
 import argparse
 
 
-pattern = re.compile(r"\s*MP_REGISTER_MODULE\((.*?),\s*(.*?)\);", flags=re.DOTALL)
+pattern = re.compile(
+    r"\s*(MP_REGISTER_MODULE|MP_REGISTER_EXTENSIBLE_MODULE)\(MP_QSTR_(.*?),\s*(.*?)\);",
+    flags=re.DOTALL,
+)
 
 
 def find_module_registrations(filename):
@@ -37,14 +49,23 @@ def generate_module_table_header(modules):
 
     # Print header file for all external modules.
     mod_defs = set()
+    extensible_mod_defs = set()
     print("// Automatically generated by makemoduledefs.py.\n")
-    for module_name, obj_module in modules:
+    for macro_name, module_name, obj_module in modules:
         mod_def = "MODULE_DEF_{}".format(module_name.upper())
-        mod_defs.add(mod_def)
+        if macro_name == "MP_REGISTER_MODULE":
+            mod_defs.add(mod_def)
+        elif macro_name == "MP_REGISTER_EXTENSIBLE_MODULE":
+            extensible_mod_defs.add(mod_def)
         if "," in obj_module:
             print(
-                "ERROR: Call to MP_REGISTER_MODULE({}, {}) should be MP_REGISTER_MODULE({}, {})\n".format(
+                "ERROR: Call to {}({}, {}) should be {}({}, {})\n".format(
-                    module_name, obj_module, module_name, obj_module.split(",")[0]
+                    macro_name,
+                    module_name,
+                    obj_module,
+                    macro_name,
+                    module_name,
+                    obj_module.split(",")[0],
                 ),
                 file=sys.stderr,
             )
@@ -53,7 +74,7 @@ def generate_module_table_header(modules):
             (
                 "extern const struct _mp_obj_module_t {obj_module};\n"
                 "#undef {mod_def}\n"
-                "#define {mod_def} {{ MP_ROM_QSTR({module_name}), MP_ROM_PTR(&{obj_module}) }},\n"
+                "#define {mod_def} {{ MP_ROM_QSTR(MP_QSTR_{module_name}), MP_ROM_PTR(&{obj_module}) }},\n"
             ).format(
                 module_name=module_name,
                 obj_module=obj_module,
@@ -68,6 +89,13 @@ def generate_module_table_header(modules):
 
     print("// MICROPY_REGISTERED_MODULES")
 
+    print("\n#define MICROPY_REGISTERED_EXTENSIBLE_MODULES \\")
+
+    for mod_def in sorted(extensible_mod_defs):
+        print("    {mod_def} \\".format(mod_def=mod_def))
+
+    print("// MICROPY_REGISTERED_EXTENSIBLE_MODULES")
+
 
 def main():
     parser = argparse.ArgumentParser()

py/makeqstrdefs.py

@@ -21,7 +21,7 @@ _MODE_QSTR = "qstr"
 # Extract MP_COMPRESSED_ROM_TEXT("") macros.  (Which come from MP_ERROR_TEXT)
 _MODE_COMPRESS = "compress"
 
-# Extract MP_REGISTER_MODULE(...) macros.
+# Extract MP_REGISTER_(EXTENSIBLE_)MODULE(...) macros.
 _MODE_MODULE = "module"
 
 # Extract MP_REGISTER_ROOT_POINTER(...) macros.
@@ -93,7 +93,7 @@ def process_file(f):
     elif args.mode == _MODE_COMPRESS:
         re_match = re.compile(r'MP_COMPRESSED_ROM_TEXT\("([^"]*)"\)')
     elif args.mode == _MODE_MODULE:
-        re_match = re.compile(r"MP_REGISTER_MODULE\(.*?,\s*.*?\);")
+        re_match = re.compile(r"MP_REGISTER_(?:EXTENSIBLE_)?MODULE\(.*?,\s*.*?\);")
     elif args.mode == _MODE_ROOT_POINTER:
         re_match = re.compile(r"MP_REGISTER_ROOT_POINTER\(.*?\);")
     output = []

py/obj.h

@@ -434,6 +434,9 @@ typedef struct _mp_rom_obj_t { mp_const_obj_t o; } mp_rom_obj_t;
 // param obj_module: mp_obj_module_t instance
 #define MP_REGISTER_MODULE(module_name, obj_module)
 
+// As above, but allow this module to be extended from the filesystem.
+#define MP_REGISTER_EXTENSIBLE_MODULE(module_name, obj_module)
+
 // Declare a root pointer (to avoid garbage collection of a global static variable).
 // param variable_declaration: a valid C variable declaration
 #define MP_REGISTER_ROOT_POINTER(variable_declaration)

py/objmodule.c

@@ -166,17 +166,54 @@ mp_obj_t mp_obj_new_module(qstr module_name) {
 // Global module table and related functions
 
 STATIC const mp_rom_map_elem_t mp_builtin_module_table[] = {
-    // builtin modules declared with MP_REGISTER_MODULE()
+    // built-in modules declared with MP_REGISTER_MODULE()
     MICROPY_REGISTERED_MODULES
 };
 MP_DEFINE_CONST_MAP(mp_builtin_module_map, mp_builtin_module_table);
 
-// Attempts to find (and initialise) a builtin, otherwise returns
+STATIC const mp_rom_map_elem_t mp_builtin_extensible_module_table[] = {
+    // built-in modules declared with MP_REGISTER_EXTENSIBLE_MODULE()
+    MICROPY_REGISTERED_EXTENSIBLE_MODULES
+};
+MP_DEFINE_CONST_MAP(mp_builtin_extensible_module_map, mp_builtin_extensible_module_table);
+
+// Attempts to find (and initialise) a built-in, otherwise returns
 // MP_OBJ_NULL.
-mp_obj_t mp_module_get_builtin(qstr module_name) {
+mp_obj_t mp_module_get_builtin(qstr module_name, bool extensible) {
-    mp_map_elem_t *elem = mp_map_lookup((mp_map_t *)&mp_builtin_module_map, MP_OBJ_NEW_QSTR(module_name), MP_MAP_LOOKUP);
+    mp_map_elem_t *elem = mp_map_lookup((mp_map_t *)(extensible ? &mp_builtin_extensible_module_map : &mp_builtin_module_map), MP_OBJ_NEW_QSTR(module_name), MP_MAP_LOOKUP);
     if (!elem) {
-        return MP_OBJ_NULL;
+        #if MICROPY_PY_SYS
+        // Special case for sys, which isn't extensible but can always be
+        // imported with the alias `usys`.
+        if (module_name == MP_QSTR_usys) {
+            return MP_OBJ_FROM_PTR(&mp_module_sys);
+        }
+        #endif
+
+        if (extensible) {
+            // At this point we've already tried non-extensible built-ins, the
+            // filesystem, and now extensible built-ins. No match, so fail
+            // the import.
+            return MP_OBJ_NULL;
+        }
+
+        // We're trying to match a non-extensible built-in (i.e. before trying
+        // the filesystem), but if the user is importing `ufoo`, _and_ `foo`
+        // is an extensible module, then allow it as a way of forcing the
+        // built-in. Essentially, this makes it as if all the extensible
+        // built-ins also had non-extensible aliases named `ufoo`. Newer code
+        // should be using sys.path to force the built-in, but this retains
+        // the old behaviour of the u-prefix being used to force a built-in
+        // import.
+        size_t module_name_len;
+        const char *module_name_str = (const char *)qstr_data(module_name, &module_name_len);
+        if (module_name_str[0] != 'u') {
+            return MP_OBJ_NULL;
+        }
+        elem = mp_map_lookup((mp_map_t *)&mp_builtin_extensible_module_map, MP_OBJ_NEW_QSTR(qstr_from_strn(module_name_str + 1, module_name_len - 1)), MP_MAP_LOOKUP);
+        if (!elem) {
+            return MP_OBJ_NULL;
+        }
     }
 
     #if MICROPY_MODULE_BUILTIN_INIT

py/objmodule.h

@@ -32,8 +32,9 @@
 #define MP_MODULE_ATTR_DELEGATION_ENTRY(ptr) { MP_ROM_QSTR(MP_QSTRnull), MP_ROM_PTR(ptr) }
 
 extern const mp_map_t mp_builtin_module_map;
+extern const mp_map_t mp_builtin_extensible_module_map;
 
-mp_obj_t mp_module_get_builtin(qstr module_name);
+mp_obj_t mp_module_get_builtin(qstr module_name, bool extensible);
 
 void mp_module_generic_attr(qstr attr, mp_obj_t *dest, const uint16_t *keys, mp_obj_t *values);
 

py/repl.c

@@ -162,8 +162,8 @@ STATIC bool test_qstr(mp_obj_t obj, qstr name) {
         return dest[0] != MP_OBJ_NULL;
     } else {
         // try builtin module
-        return mp_map_lookup((mp_map_t *)&mp_builtin_module_map,
+        return mp_map_lookup((mp_map_t *)&mp_builtin_module_map, MP_OBJ_NEW_QSTR(name), MP_MAP_LOOKUP) ||
-            MP_OBJ_NEW_QSTR(name), MP_MAP_LOOKUP);
+               mp_map_lookup((mp_map_t *)&mp_builtin_extensible_module_map, MP_OBJ_NEW_QSTR(name), MP_MAP_LOOKUP);
     }
 }