py/builtinimport: Support relative import in custom __import__ callback.

The globals need to be forwarded from the caller's context.

Signed-off-by: Damien George <damien@micropython.org>

4 files changed, 42 insertions, 6 deletions

diff --git a/py/builtinimport.c b/py/builtinimport.c
index ff894f69d..2c7d79668 100644
--- a/py/builtinimport.c
+++ b/py/builtinimport.c
@@ -267,7 +267,7 @@ static void do_load(mp_module_context_t *module_obj, vstr_t *file) {
 
 // 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) {
+static void evaluate_relative_import(mp_int_t level, const char **module_name, size_t *module_name_len, mp_obj_t globals) {
     // 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.
@@ -276,7 +276,7 @@ static void evaluate_relative_import(mp_int_t level, const char **module_name, s
     // 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__));
+    mp_obj_t current_module_name_obj = mp_obj_dict_get(globals, MP_OBJ_NEW_QSTR(MP_QSTR___name__));
     assert(current_module_name_obj != MP_OBJ_NULL);
 
     #if MICROPY_MODULE_OVERRIDE_MAIN_IMPORT && MICROPY_CPYTHON_COMPAT
@@ -284,12 +284,12 @@ static void evaluate_relative_import(mp_int_t level, const char **module_name, s
         // 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__));
+        current_module_name_obj = mp_obj_dict_get(globals, 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);
+    bool is_pkg = mp_map_lookup(mp_obj_dict_get_map(globals), MP_OBJ_NEW_QSTR(MP_QSTR___path__), MP_MAP_LOOKUP);
 
     #if DEBUG_PRINT
     DEBUG_printf("Current module/package: ");
@@ -569,10 +569,19 @@ mp_obj_t mp_builtin___import___default(size_t n_args, const mp_obj_t *args) {
     const char *module_name = mp_obj_str_get_data(module_name_obj, &module_name_len);
 
     if (level != 0) {
+        // This is the dict with all global symbols.
+        mp_obj_t globals = MP_OBJ_FROM_PTR(mp_globals_get());
+        if (n_args >= 2 && args[1] != mp_const_none) {
+            globals = args[1];
+            if (!mp_obj_is_type(globals, &mp_type_dict)) {
+                mp_raise_TypeError(NULL);
+            }
+        }
+
         // Turn "foo.bar" with level=3 into "<current module 3 components>.foo.bar".
         // Current module name is extracted from globals().__name__.
-        evaluate_relative_import(level, &module_name, &module_name_len);
         // module_name is now an absolute module path.
+        evaluate_relative_import(level, &module_name, &module_name_len, globals);
     }
 
     if (module_name_len == 0) {
diff --git a/py/runtime.c b/py/runtime.c
index a84e22760..ddc1ef7d3 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -1522,7 +1522,7 @@ mp_obj_t mp_import_name(qstr name, mp_obj_t fromlist, mp_obj_t level) {
     // build args array
     mp_obj_t args[5];
     args[0] = MP_OBJ_NEW_QSTR(name);
-    args[1] = mp_const_none; // TODO should be globals
+    args[1] = MP_OBJ_FROM_PTR(mp_globals_get()); // globals of the current context
     args[2] = mp_const_none; // TODO should be locals
     args[3] = fromlist;
     args[4] = level;
diff --git a/tests/import/builtin_import.py b/tests/import/builtin_import.py
index 734498d1b..364b0bae9 100644
--- a/tests/import/builtin_import.py
+++ b/tests/import/builtin_import.py
@@ -20,3 +20,12 @@ try:
     __import__("xyz", None, None, None, -1)
 except ValueError:
     print("ValueError")
+
+# globals is not checked for level=0
+__import__("builtins", "globals")
+
+# globals must be a dict (or None) for level>0
+try:
+    __import__("builtins", "globals", None, None, 1)
+except TypeError:
+    print("TypeError")
diff --git a/tests/import/import_override2.py b/tests/import/import_override2.py
new file mode 100644
index 000000000..25aac44fe
--- /dev/null
+++ b/tests/import/import_override2.py
@@ -0,0 +1,18 @@
+# test overriding __import__ combined with importing from the filesystem
+
+
+def custom_import(name, globals, locals, fromlist, level):
+    if level > 0:
+        print("import", name, fromlist, level)
+    return orig_import(name, globals, locals, fromlist, level)
+
+
+orig_import = __import__
+try:
+    __import__("builtins").__import__ = custom_import
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+# import calls __import__ behind the scenes
+import pkg7.subpkg1.subpkg2.mod3