Merge pull request #148 from tjhance/fix-import

(more) correct semantics for import

src/analysis/type_analysis.cpp

@@ -365,7 +365,11 @@ private:
             case AST_LangPrimitive::LANDINGPAD:
             case AST_LangPrimitive::GET_ITER:
             case AST_LangPrimitive::IMPORT_FROM:
+            case AST_LangPrimitive::IMPORT_STAR:
+            case AST_LangPrimitive::IMPORT_NAME:
                 return UNKNOWN;
+            case AST_LangPrimitive::NONE:
+                return NONE;
             default:
                 RELEASE_ASSERT(0, "%d", node->opcode);
         }

src/codegen/irgen/irgenerator.cpp

@@ -431,6 +431,52 @@ private:
                 converted_module->decvref(emitter);
                 return v;
             }
+            case AST_LangPrimitive::IMPORT_STAR: {
+                assert(node->args.size() == 1);
+                assert(node->args[0]->type == AST_TYPE::Name);
+
+                RELEASE_ASSERT(irstate->getSourceInfo()->ast->type == AST_TYPE::Module,
+                               "import * not supported in functions");
+
+                CompilerVariable* module = evalExpr(node->args[0], exc_info);
+                ConcreteCompilerVariable* converted_module = module->makeConverted(emitter, module->getBoxType());
+                module->decvref(emitter);
+
+                llvm::Value* r = emitter.createCall2(exc_info, g.funcs.importStar, converted_module->getValue(),
+                                                     embedConstantPtr(irstate->getSourceInfo()->parent_module,
+                                                                      g.llvm_module_type_ptr)).getInstruction();
+                CompilerVariable* v = new ConcreteCompilerVariable(UNKNOWN, r, true);
+
+                converted_module->decvref(emitter);
+                return v;
+            }
+            case AST_LangPrimitive::IMPORT_NAME: {
+                assert(node->args.size() == 3);
+                assert(node->args[0]->type == AST_TYPE::Num);
+                assert(static_cast<AST_Num*>(node->args[0])->num_type == AST_Num::INT);
+                assert(node->args[2]->type == AST_TYPE::Str);
+
+                int level = static_cast<AST_Num*>(node->args[0])->n_int;
+
+                // TODO this could be a constant Box* too
+                CompilerVariable* froms = evalExpr(node->args[1], exc_info);
+                ConcreteCompilerVariable* converted_froms = froms->makeConverted(emitter, froms->getBoxType());
+                froms->decvref(emitter);
+
+                const std::string& module_name = static_cast<AST_Str*>(node->args[2])->s;
+
+                llvm::Value* imported
+                    = emitter.createCall3(exc_info, g.funcs.import, getConstantInt(level, g.i32),
+                                          converted_froms->getValue(),
+                                          embedConstantPtr(&module_name, g.llvm_str_type_ptr)).getInstruction();
+                ConcreteCompilerVariable* v = new ConcreteCompilerVariable(UNKNOWN, imported, true);
+
+                converted_froms->decvref(emitter);
+                return v;
+            }
+            case AST_LangPrimitive::NONE: {
+                return getNone();
+            }
             default:
                 RELEASE_ASSERT(0, "%d", node->opcode);
         }
@@ -1611,54 +1657,6 @@ private:
         func->decvref(emitter);
     }
 
-    void doImport(AST_Import* node, ExcInfo exc_info) {
-        if (state == PARTIAL)
-            return;
-
-        for (int i = 0; i < node->names.size(); i++) {
-            AST_alias* alias = node->names[i];
-
-            const std::string& modname = alias->name;
-            const std::string& asname = alias->asname.size() ? alias->asname : alias->name;
-
-            llvm::Value* imported
-                = emitter.createCall(exc_info, g.funcs.import, embedConstantPtr(&modname, g.llvm_str_type_ptr))
-                      .getInstruction();
-            ConcreteCompilerVariable* v = new ConcreteCompilerVariable(UNKNOWN, imported, true);
-            _doSet(asname, v, exc_info);
-            v->decvref(emitter);
-        }
-    }
-
-    void doImportFrom(AST_ImportFrom* node, ExcInfo exc_info) {
-        if (state == PARTIAL)
-            return;
-
-        assert(node->level == 0);
-
-        llvm::Value* imported_v
-            = emitter.createCall(exc_info, g.funcs.import, embedConstantPtr(&node->module, g.llvm_str_type_ptr))
-                  .getInstruction();
-        ConcreteCompilerVariable* module = new ConcreteCompilerVariable(typeFromClass(module_cls), imported_v, true);
-
-        // ImportFrom statements should only get generated for "import *"; all other import from's
-        // should get translated to IMPORT_FROM bytecodes.
-        for (int i = 0; i < node->names.size(); i++) {
-            AST_alias* alias = node->names[i];
-
-            assert(alias->name == "*");
-            assert(alias->asname == "");
-
-            RELEASE_ASSERT(irstate->getSourceInfo()->ast->type == AST_TYPE::Module,
-                           "import * not supported in functions");
-            emitter.createCall2(exc_info, g.funcs.importStar, imported_v,
-                                embedConstantPtr(irstate->getSourceInfo()->parent_module, g.llvm_module_type_ptr));
-            continue;
-        }
-
-        module->decvref(emitter);
-    }
-
     void doPrint(AST_Print* node, ExcInfo exc_info) {
         if (state == PARTIAL)
             return;
@@ -2043,12 +2041,12 @@ private:
             // case AST_TYPE::If:
             // doIf(ast_cast<AST_If>(node));
             // break;
-            case AST_TYPE::Import:
-                doImport(ast_cast<AST_Import>(node), exc_info);
-                break;
-            case AST_TYPE::ImportFrom:
-                doImportFrom(ast_cast<AST_ImportFrom>(node), exc_info);
-                break;
+            // case AST_TYPE::Import:
+            //     doImport(ast_cast<AST_Import>(node), exc_info);
+            //     break;
+            // case AST_TYPE::ImportFrom:
+            //     doImportFrom(ast_cast<AST_ImportFrom>(node), exc_info);
+            //     break;
             case AST_TYPE::Global:
                 // Should have been handled already
                 break;

src/core/ast.cpp

@@ -1371,6 +1371,15 @@ bool PrintVisitor::visit_langprimitive(AST_LangPrimitive* node) {
         case AST_LangPrimitive::IMPORT_FROM:
             printf("IMPORT_FROM");
             break;
+        case AST_LangPrimitive::IMPORT_NAME:
+            printf("IMPORT_NAME");
+            break;
+        case AST_LangPrimitive::IMPORT_STAR:
+            printf("IMPORT_STAR");
+            break;
+        case AST_LangPrimitive::NONE:
+            printf("NONE");
+            break;
         default:
             RELEASE_ASSERT(0, "%d", node->opcode);
     }

src/core/ast.h

@@ -256,6 +256,9 @@ public:
 
     AST_Attribute() : AST_expr(AST_TYPE::Attribute) {}
 
+    AST_Attribute(AST_expr* value, AST_TYPE::AST_TYPE ctx_type, const std::string& attr)
+        : AST_expr(AST_TYPE::Attribute), value(value), ctx_type(ctx_type), attr(attr) {}
+
     static const AST_TYPE::AST_TYPE TYPE = AST_TYPE::Attribute;
 };
 
@@ -923,6 +926,9 @@ public:
         LOCALS,
         GET_ITER,
         IMPORT_FROM,
+        IMPORT_NAME,
+        IMPORT_STAR,
+        NONE,
     } opcode;
     std::vector<AST_expr*> args;
 

src/core/cfg.cpp

@@ -1083,18 +1083,50 @@ public:
         return true;
     }
 
+    static std::string getTopModule(const std::string& full_name) {
+        size_t period_index = full_name.find('.');
+        if (period_index == std::string::npos) {
+            return full_name;
+        } else {
+            return full_name.substr(0, period_index);
+        }
+    }
+
     virtual bool visit_import(AST_Import* node) {
         for (AST_alias* a : node->names) {
-            AST_Import* remapped = new AST_Import();
-            remapped->lineno = node->lineno;
-            remapped->col_offset = node->col_offset;
+            AST_LangPrimitive* import = new AST_LangPrimitive(AST_LangPrimitive::IMPORT_NAME);
+            import->args.push_back(new AST_Num());
+            static_cast<AST_Num*>(import->args[0])->num_type = AST_Num::INT;
+            static_cast<AST_Num*>(import->args[0])->n_int = -1;
+            import->args.push_back(new AST_LangPrimitive(AST_LangPrimitive::NONE));
+            import->args.push_back(new AST_Str(a->name));
 
             std::string tmpname = nodeName(a);
-            AST_alias* remapped_alias = new AST_alias(a->name, tmpname);
-            remapped->names.push_back(remapped_alias);
+            pushAssign(tmpname, import);
 
-            push_back(remapped);
-            pushAssign(a->asname.size() ? a->asname : a->name, makeName(tmpname, AST_TYPE::Load));
+            if (a->asname.size() == 0) {
+                // No asname, so load the top-level module into the name
+                // (e.g., for `import os.path`, loads the os module into `os`)
+                pushAssign(getTopModule(a->name), makeName(tmpname, AST_TYPE::Load));
+            } else {
+                // If there is an asname, get the bottom-level module by
+                // getting the attributes and load it into asname.
+                int l = 0;
+                do {
+                    int r = a->name.find('.', l);
+                    if (r == std::string::npos) {
+                        r = a->name.size();
+                    }
+                    if (l == 0) {
+                        l = r + 1;
+                        continue;
+                    }
+                    pushAssign(tmpname, new AST_Attribute(makeName(tmpname, AST_TYPE::Load), AST_TYPE::Load,
+                                                          a->name.substr(l, r)));
+                    l = r + 1;
+                } while (l < a->name.size());
+                pushAssign(a->asname, makeName(tmpname, AST_TYPE::Load));
+            }
         }
 
         return true;
@@ -1103,27 +1135,32 @@ public:
     virtual bool visit_importfrom(AST_ImportFrom* node) {
         RELEASE_ASSERT(node->level == 0, "");
 
-        AST_Import* import = new AST_Import();
-        import->lineno = node->lineno;
-        import->col_offset = node->col_offset;
+        AST_LangPrimitive* import = new AST_LangPrimitive(AST_LangPrimitive::IMPORT_NAME);
+        import->args.push_back(new AST_Num());
+        static_cast<AST_Num*>(import->args[0])->num_type = AST_Num::INT;
+        static_cast<AST_Num*>(import->args[0])->n_int = node->level;
+        import->args.push_back(new AST_Tuple());
+        static_cast<AST_Tuple*>(import->args[1])->ctx_type = AST_TYPE::Load;
+        for (int i = 0; i < node->names.size(); i++) {
+            static_cast<AST_Tuple*>(import->args[1])->elts.push_back(new AST_Str(node->names[i]->name));
+        }
+        import->args.push_back(new AST_Str(node->module));
 
         std::string tmp_module_name = nodeName(node);
-        AST_alias* tmp_alias = new AST_alias(node->module, tmp_module_name);
-        import->names.push_back(tmp_alias);
-        push_back(import);
+        pushAssign(tmp_module_name, import);
 
         for (AST_alias* a : node->names) {
             if (a->name == "*") {
-                assert(a->asname.size() == 0);
 
-                AST_ImportFrom* remapped = new AST_ImportFrom();
-                remapped->lineno = node->lineno;
-                remapped->col_offset = node->col_offset;
-                remapped->module = node->module;
-                remapped->level = node->level;
-                remapped->names.push_back(a);
+                AST_LangPrimitive* import_star = new AST_LangPrimitive(AST_LangPrimitive::IMPORT_STAR);
+                import_star->lineno = node->lineno;
+                import_star->col_offset = node->col_offset;
+                import_star->args.push_back(makeName(tmp_module_name, AST_TYPE::Load));
+
+                AST_Expr* import_star_expr = new AST_Expr();
+                import_star_expr->value = import_star;
 
-                push_back(remapped);
+                push_back(import_star_expr);
             } else {
                 AST_LangPrimitive* import_from = new AST_LangPrimitive(AST_LangPrimitive::IMPORT_FROM);
                 import_from->lineno = node->lineno;

src/runtime/capi.cpp

@@ -627,7 +627,7 @@ extern "C" PyObject* PyImport_Import(PyObject* module_name) {
     RELEASE_ASSERT(module_name->cls == str_cls, "");
 
     try {
-        return import(&static_cast<BoxedString*>(module_name)->s);
+        return import(-1, None, &static_cast<BoxedString*>(module_name)->s);
     } catch (Box* e) {
         Py_FatalError("unimplemented");
     }

src/runtime/import.cpp

@@ -38,24 +38,7 @@ BoxedModule* createAndRunModule(const std::string& name, const std::string& fn)
 #define LLVM_SYS_FS_EXISTS_CODE_OKAY(code) (!(code))
 #endif
 
-extern "C" Box* import(const std::string* name) {
-    assert(name);
-
-    static StatCounter slowpath_import("slowpath_import");
-    slowpath_import.log();
-
-    Box* parent_module = NULL;
-    size_t periodIndex = name->rfind('.');
-    if (periodIndex != std::string::npos) {
-        std::string parent_name(*name, 0, periodIndex);
-        parent_module = import(&parent_name);
-    }
-
-    BoxedDict* sys_modules = getSysModulesDict();
-    Box* s = boxStringPtr(name);
-    if (sys_modules->d.find(s) != sys_modules->d.end())
-        return sys_modules->d[s];
-
+static Box* importSub(const std::string* name, Box* parent_module) {
     BoxedList* path_list;
     if (parent_module == NULL) {
         path_list = getSysPath();
@@ -103,4 +86,45 @@ extern "C" Box* import(const std::string* name) {
 
     raiseExcHelper(ImportError, "No module named %s", name->c_str());
 }
+
+static Box* import(const std::string* name, bool return_first) {
+    assert(name);
+    assert(name->size() > 0);
+
+    static StatCounter slowpath_import("slowpath_import");
+    slowpath_import.log();
+
+    BoxedDict* sys_modules = getSysModulesDict();
+
+    size_t l = 0, r;
+    Box* last_module = NULL;
+    Box* first_module = NULL;
+    while (l < name->size()) {
+        size_t r = name->find('.', l);
+        if (r == std::string::npos) {
+            r = name->size();
+        }
+
+        std::string prefix_name = std::string(*name, 0, r);
+        Box* s = boxStringPtr(&prefix_name);
+        if (sys_modules->d.find(s) != sys_modules->d.end()) {
+            last_module = sys_modules->d[s];
+        } else {
+            std::string small_name = std::string(*name, l, r - l);
+            last_module = importSub(&small_name, last_module);
+        }
+
+        if (l == 0) {
+            first_module = last_module;
+        }
+
+        l = r + 1;
+    }
+
+    return return_first ? first_module : last_module;
+}
+
+extern "C" Box* import(int level, Box* from_imports, const std::string* module_name) {
+    return import(module_name, from_imports == None);
+}
 }

src/runtime/import.h

@@ -19,7 +19,7 @@
 
 namespace pyston {
 
-extern "C" Box* import(const std::string* name);
+extern "C" Box* import(int level, Box* from_imports, const std::string* module_name);
 }
 
 #endif

src/runtime/objmodel.cpp

@@ -3674,7 +3674,7 @@ extern "C" Box* importFrom(Box* _m, const std::string* name) {
     raiseExcHelper(ImportError, "cannot import name %s", name->c_str());
 }
 
-extern "C" void importStar(Box* _from_module, BoxedModule* to_module) {
+extern "C" Box* importStar(Box* _from_module, BoxedModule* to_module) {
     assert(_from_module->cls == module_cls);
     BoxedModule* from_module = static_cast<BoxedModule*>(_from_module);
 
@@ -3710,7 +3710,7 @@ extern "C" void importStar(Box* _from_module, BoxedModule* to_module) {
 
             to_module->setattr(casted_attr_name->s, attr_value, NULL);
         }
-        return;
+        return None;
     }
 
     HCAttrs* module_attrs = from_module->getAttrsPtr();
@@ -3720,5 +3720,7 @@ extern "C" void importStar(Box* _from_module, BoxedModule* to_module) {
 
         to_module->setattr(p.first, module_attrs->attr_list->attrs[p.second], NULL);
     }
+
+    return None;
 }
 }

src/runtime/objmodel.h

@@ -76,7 +76,7 @@ extern "C" void delitem(Box* target, Box* slice);
 extern "C" Box* getclsattr(Box* obj, const char* attr);
 extern "C" Box* unaryop(Box* operand, int op_type);
 extern "C" Box* importFrom(Box* obj, const std::string* attr);
-extern "C" void importStar(Box* from_module, BoxedModule* to_module);
+extern "C" Box* importStar(Box* from_module, BoxedModule* to_module);
 extern "C" Box** unpackIntoArray(Box* obj, int64_t expected_size);
 extern "C" void assertNameDefined(bool b, const char* name, BoxedClass* exc_cls, bool local_var_msg);
 extern "C" void assertFail(BoxedModule* inModule, Box* msg);

test/tests/os_path.py

 # allow-warning: converting unicode literal to str
 
 import os.path
+
+print type(os)
+print type(os.path)
+
+import os.path as foo
+print type(foo)
+
+from os import path
+print type(path)
+
+from os import path as moo
+print type(moo)
+
+from os.path import abspath
+print type(abspath)
+
+from os.path import abspath as llama
+print type(llama)
+
+from os.path import *
+print type(commonprefix)