Merge pull request #373 from iabdalkader/module_register

Add mp_obj_module_register

3 files changed, 8 insertions, 3 deletions

diff --git a/py/obj.h b/py/obj.h
index 49e92795c6..dcb43c0748 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -456,6 +456,7 @@ typedef struct _mp_obj_module_t {
 extern const mp_obj_type_t mp_type_module;
 mp_obj_t mp_obj_new_module(qstr module_name);
 mp_obj_t mp_obj_module_get(qstr module_name);
+void mp_obj_module_register(qstr qstr, mp_obj_t module); //use for loading statically allocated modules
 struct _mp_map_t *mp_obj_module_get_globals(mp_obj_t self_in);
 
 // staticmethod and classmethod types; defined here so we can make const versions
diff --git a/py/objmodule.c b/py/objmodule.c
index e0778b4cf2..21fec9e866 100644
--- a/py/objmodule.c
+++ b/py/objmodule.c
@@ -82,6 +82,11 @@ mp_obj_t mp_obj_module_get(qstr module_name) {
     return MP_OBJ_NULL;
 }
 
+void mp_obj_module_register(qstr qstr, mp_obj_t module)
+{
+    mp_map_lookup(rt_loaded_modules_get(), MP_OBJ_NEW_QSTR(qstr), MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = module;
+}
+
 mp_map_t *mp_obj_module_get_globals(mp_obj_t self_in) {
     assert(MP_OBJ_IS_TYPE(self_in, &mp_type_module));
     mp_obj_module_t *self = self_in;
diff --git a/stm/main.c b/stm/main.c
index 8279a9b7f4..6ad3111276 100644
--- a/stm/main.c
+++ b/stm/main.c
@@ -272,9 +272,8 @@ soft_reset:
     rt_store_name(MP_QSTR_help, rt_make_function_n(0, pyb_help));
     rt_store_name(MP_QSTR_open, rt_make_function_n(2, pyb_io_open));
 
-    // we pre-import the pyb module
-    // probably shouldn't do this, so we are compatible with CPython
-    rt_store_name(MP_QSTR_pyb, (mp_obj_t)&pyb_module);
+    // load the pyb module
+    mp_obj_module_register(MP_QSTR_pyb, (mp_obj_t)&pyb_module);
 
     // check if user switch held (initiates reset of filesystem)
     bool reset_filesystem = false;