import: Implement "from pkg.mod import sym" syntax properly.

http://docs.python.org/3.3/library/functions.html#__import__ :
"When the name variable is of the form package.module, normally, the
top-level package (the name up till the first dot) is returned, not
the module named by name. However, when a non-empty fromlist argument
is given, the module named by name is returned."

2 files changed, 34 insertions, 3 deletions

diff --git a/py/builtinimport.c b/py/builtinimport.c
index 0a730b031c..05b8eead57 100644
--- a/py/builtinimport.c
+++ b/py/builtinimport.c
@@ -132,13 +132,18 @@ void do_load(mp_obj_t module_obj, vstr_t *file) {
 mp_obj_t mp_builtin___import__(int n_args, mp_obj_t *args) {
     /*
     printf("import:\n");
-    for (int i = 0; i < n; i++) {
+    for (int i = 0; i < n_args; i++) {
         printf("  ");
-        mp_obj_print(args[i]);
+        mp_obj_print(args[i], PRINT_REPR);
         printf("\n");
     }
     */
 
+    mp_obj_t fromtuple = mp_const_none;
+    if (n_args >= 4) {
+        fromtuple = args[3];
+    }
+
     uint mod_len;
     const char *mod_str = (const char*)mp_obj_str_get_data(args[0], &mod_len);
 
@@ -150,8 +155,11 @@ mp_obj_t mp_builtin___import__(int n_args, mp_obj_t *args) {
         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
-        // TODO: subject to fromlist arg
         qstr pkg_name = qstr_from_strn(mod_str, p - mod_str);
         return mp_obj_module_get(pkg_name);
     }
@@ -227,6 +235,11 @@ mp_obj_t mp_builtin___import__(int n_args, mp_obj_t *args) {
         assert(0);
     }
 
+    // If fromlist is not empty, return leaf module
+    if (fromtuple != mp_const_none) {
+        return module_obj;
+    }
+    // Otherwise, we need to return top-level package
     return top_module_obj;
 }
 
diff --git a/tests/basics/import-pkg2.py b/tests/basics/import-pkg2.py
new file mode 100644
index 0000000000..2e9f34121b
--- /dev/null
+++ b/tests/basics/import-pkg2.py
@@ -0,0 +1,18 @@
+from pkg.mod import foo
+
+try:
+    pkg
+except NameError:
+    print("NameError")
+try:
+    pkg.mod
+except NameError:
+    print("NameError")
+print(foo())
+
+# Import few times, must be same module objects
+mod_1 = __import__("pkg.mod", None, None, ("foo",))
+mod_2 = __import__("pkg.mod", None, None, ("foo",))
+print(mod_1 is mod_2)
+print(mod_1.foo is mod_2.foo)
+print(foo is mod_1.foo)