Support import statements with custom globals

imports apparently have to look at the globals scope in order
to determine how to handle relative imports and package references.

Sqlalchemy for some reason does this, and once we have this sqlalchemy
imports.

src/codegen/ast_interpreter.cpp

@@ -342,12 +342,7 @@ Value ASTInterpreter::doBinOp(Box* left, Box* right, int op, BinExpType exp_type
 void ASTInterpreter::doStore(InternedString name, Value value) {
     ScopeInfo::VarScopeType vst = scope_info->getScopeTypeOfName(name);
     if (vst == ScopeInfo::VarScopeType::GLOBAL) {
-        if (globals->cls == module_cls) {
-            setattr(static_cast<BoxedModule*>(globals), name.c_str(), value.o);
-        } else {
-            assert(globals->cls == dict_cls);
-            static_cast<BoxedDict*>(globals)->d[boxString(name.str())] = value.o;
-        }
+        setGlobal(globals, name, value.o);
     } else if (vst == ScopeInfo::VarScopeType::NAME) {
         assert(frame_info.boxedLocals != NULL);
         // TODO should probably pre-box the names when it's a scope that usesNameLookup
@@ -879,7 +874,6 @@ Value ASTInterpreter::visit_delete(AST_Delete* node) {
 
                 ScopeInfo::VarScopeType vst = scope_info->getScopeTypeOfName(target->id);
                 if (vst == ScopeInfo::VarScopeType::GLOBAL) {
-                    // Can't use delattr since the errors are different:
                     delGlobal(globals, &target->id.str());
                     continue;
                 } else if (vst == ScopeInfo::VarScopeType::NAME) {

src/codegen/unwinding.cpp

@@ -599,6 +599,8 @@ CompiledFunction* getTopCompiledFunction() {
 
 Box* getGlobals() {
     auto it = getTopPythonFrame();
+    if (!it)
+        return NULL;
     return it->getGlobals();
 }
 

src/core/stringpool.h

@@ -80,6 +80,8 @@ public:
         return *this->_str < *rhs._str;
     }
 
+    operator llvm::StringRef() { return llvm::StringRef(*_str); }
+
     friend class InternedStringPool;
     friend struct std::hash<InternedString>;
     friend struct std::less<InternedString>;

src/jit.cpp

@@ -286,7 +286,10 @@ static int main(int argc, char** argv) {
 
             llvm::sys::path::append(path, fn);
             llvm::sys::path::remove_filename(path);
-            prependToSysPath(path.str());
+            char* real_path
+                = realpath(path.str().str().c_str(), NULL); // inefficient way of null-terminating the string
+            prependToSysPath(real_path);
+            free(real_path);
 
             main_module = createModule("__main__", fn);
             try {

src/runtime/builtin_modules/builtins.cpp

@@ -388,6 +388,10 @@ Box* bltinImport(Box* name, Box* globals, Box* locals, Box** args) {
     Box* fromlist = args[0];
     Box* level = args[1];
 
+    // __import__ takes a 'locals' argument, but it doesn't get used in CPython.
+    // Well, it gets passed to PyImport_ImportModuleLevel() and then import_module_level(),
+    // which ignores it.  So we don't even pass it through.
+
     name = coerceUnicodeToStr(name);
 
     if (name->cls != str_cls) {

src/runtime/import.cpp

@@ -21,6 +21,7 @@
 
 #include "codegen/irgen/hooks.h"
 #include "codegen/parser.h"
+#include "codegen/unwinding.h"
 #include "runtime/capi.h"
 #include "runtime/objmodel.h"
 
@@ -267,7 +268,7 @@ static Box* getParent(Box* globals, int level, std::string& buf) {
     if (globals == NULL || globals == None || level == 0)
         return None;
 
-    BoxedString* pkgname = static_cast<BoxedString*>(getattrInternal(globals, package_str, NULL));
+    BoxedString* pkgname = static_cast<BoxedString*>(getFromGlobals(globals, package_str));
     if (pkgname != NULL && pkgname != None) {
         /* __package__ is set, so use it */
         if (pkgname->cls != str_cls) {
@@ -286,17 +287,11 @@ static Box* getParent(Box* globals, int level, std::string& buf) {
         buf += pkgname->s;
     } else {
         /* __package__ not set, so figure it out and set it */
-        BoxedString* modname = static_cast<BoxedString*>(getattrInternal(globals, name_str, NULL));
+        BoxedString* modname = static_cast<BoxedString*>(getFromGlobals(globals, name_str));
         if (modname == NULL || modname->cls != str_cls)
             return None;
 
-        Box* modpath = NULL;
-        try {
-            modpath = getattrInternal(globals, path_str, NULL);
-        } catch (ExcInfo e) {
-            if (!e.matches(AttributeError))
-                raiseRaw(e);
-        }
+        Box* modpath = getFromGlobals(globals, path_str);
 
         if (modpath != NULL) {
             /* __path__ is set, so modname is already the package name */
@@ -304,7 +299,7 @@ static Box* getParent(Box* globals, int level, std::string& buf) {
                 raiseExcHelper(ValueError, "Module name too long");
             }
             buf += modname->s;
-            globals->setattr(package_str, modname, NULL);
+            setGlobal(globals, package_str, modname);
         } else {
             /* Normal module, so work out the package name if any */
             size_t lastdot = modname->s.rfind('.');
@@ -312,7 +307,7 @@ static Box* getParent(Box* globals, int level, std::string& buf) {
                 raiseExcHelper(ValueError, "Attempted relative import in non-package");
             }
             if (lastdot == std::string::npos) {
-                globals->setattr(package_str, None, NULL);
+                setGlobal(globals, package_str, None);
                 return None;
             }
             if (lastdot >= PATH_MAX) {
@@ -320,7 +315,7 @@ static Box* getParent(Box* globals, int level, std::string& buf) {
             }
 
             buf = std::string(modname->s, 0, lastdot);
-            globals->setattr(package_str, boxStringPtr(&buf), NULL);
+            setGlobal(globals, package_str, boxStringPtr(&buf));
         }
     }
 
@@ -471,7 +466,6 @@ static bool loadNext(Box* mod, Box* altmod, std::string& name, std::string& buf,
 
 static void ensureFromlist(Box* module, Box* fromlist, std::string& buf, bool recursive);
 Box* importModuleLevel(const std::string& name, Box* globals, Box* from_imports, int level) {
-    RELEASE_ASSERT(!globals || globals == None || isSubclass(globals->cls, module_cls), "");
     bool return_first = from_imports == None;
 
     static StatCounter slowpath_import("slowpath_import");
@@ -665,7 +659,7 @@ Box* nullImporterFindModule(Box* self, Box* fullname, Box* path) {
 
 extern "C" Box* import(int level, Box* from_imports, const std::string* module_name) {
     std::string _module_name(*module_name);
-    return importModuleLevel(_module_name, getCurrentModule(), from_imports, level);
+    return importModuleLevel(_module_name, getGlobals(), from_imports, level);
 }
 
 Box* impFindModule(Box* _name, BoxedList* path) {

src/runtime/objmodel.cpp

@@ -4562,6 +4562,40 @@ extern "C" Box* getGlobal(Box* globals, const std::string* name) {
     raiseExcHelper(NameError, "global name '%s' is not defined", name->c_str());
 }
 
+Box* getFromGlobals(Box* globals, llvm::StringRef name) {
+    if (globals->cls == attrwrapper_cls) {
+        globals = unwrapAttrWrapper(globals);
+        RELEASE_ASSERT(globals->cls == module_cls, "%s", globals->cls->tp_name);
+    }
+
+    if (globals->cls == module_cls) {
+        return globals->getattr(name);
+    } else if (globals->cls == dict_cls) {
+        auto d = static_cast<BoxedDict*>(globals)->d;
+        auto name_str = boxString(name.str());
+        auto it = d.find(name_str);
+        if (it != d.end())
+            return it->second;
+        return NULL;
+    } else {
+        RELEASE_ASSERT(0, "%s", globals->cls->tp_name);
+    }
+}
+
+void setGlobal(Box* globals, llvm::StringRef name, Box* value) {
+    if (globals->cls == attrwrapper_cls) {
+        globals = unwrapAttrWrapper(globals);
+        RELEASE_ASSERT(globals->cls == module_cls, "%s", globals->cls->tp_name);
+    }
+
+    if (globals->cls == module_cls) {
+        setattr(static_cast<BoxedModule*>(globals), name.data(), value);
+    } else {
+        RELEASE_ASSERT(globals->cls == dict_cls, "%s", globals->cls->tp_name);
+        static_cast<BoxedDict*>(globals)->d[boxString(name)] = value;
+    }
+}
+
 extern "C" Box* importFrom(Box* _m, const std::string* name) {
     STAT_TIMER(t0, "us_timer_importFrom");
 

src/runtime/objmodel.h

@@ -183,7 +183,13 @@ inline std::tuple<Box*, Box*, Box*, Box**> getTupleFromArgsArray(Box** args, int
 
 // The `globals` argument can be either a BoxedModule or a BoxedDict
 
+// Corresponds to a name lookup with GLOBAL scope.  Checks the passed globals object, then the builtins,
+// and if not found raises an exception.
 extern "C" Box* getGlobal(Box* globals, const std::string* name);
+// Checks for the name just in the passed globals object, and returns NULL if it is not found.
+// This includes if the globals object defined a custom __getattr__ method that threw an AttributeError.
+Box* getFromGlobals(Box* globals, llvm::StringRef name);
+void setGlobal(Box* globals, llvm::StringRef name, Box* value);
 extern "C" void delGlobal(Box* globals, const std::string* name);
 
 extern "C" void boxedLocalsSet(Box* boxedLocals, const char* attr, Box* val);

test/integration/virtualenv_test.py

@@ -21,9 +21,10 @@ set -e
 set -ux
 python -c 'import __future__'
 python -c 'import sys; print sys.executable'
-pip install bcrypt==1.1.0 python-gflags==2.0
+pip install bcrypt==1.1.0 python-gflags==2.0 sqlalchemy==1.0.0
 python -c 'import bcrypt; assert bcrypt.__version__ == "1.1.0"; assert bcrypt.hashpw("password1", "$2a$12$0123456789012345678901").endswith("I1hdtg4K"); print "bcrypt seems to work"'
 python -c 'import gflags; print "gflags imports"'
+python -c 'import sqlalchemy; print "sqlalchemy imports"'
 """.strip()
 
 # print sh_script

test/tests/import_globals.py

+import test_package.import_target
+import import_target
+print
+
+s = """
+import import_target
+print import_target.__file__.replace(".pyc", ".py")
+"""
+
+def test(g):
+    print "Testing with globals:", "None" if g is None else ("globals" if g is globals() else sorted(g.items()))
+    print __import__("import_target", g).__file__.replace(".pyc", ".py")
+    exec s in g
+    print "Resulting globals:", "None" if g is None else ("globals" if g is globals() else sorted(g.keys()))
+    print
+
+test(None)
+test(globals())
+
+test({"__package__":"test_package"})
+
+test({"__name__":"test_package.foo"})
+test({"__name__":"test_package"})
+test({"__name__":"test_package", "__path__":"foo"})