Switch uses of getCurrentModule to a new getGlobals()

The callers were assuming that the module represents the globals,
which is not always true.

We can currently fetch globals from the interpreter, but not the jit.
We don't allow custom globals through the jit, so for now we can just
return the module; I think we should be able to support the jit by adding
the globals as a stackmap arg.

src/codegen/ast_interpreter.cpp

@@ -150,6 +150,11 @@ public:
         return current_inst;
     }
 
+    Box* getGlobals() {
+        assert(globals);
+        return globals;
+    }
+
     CompiledFunction* getCF() { return compiled_func; }
     FrameInfo* getFrameInfo() { return &frame_info; }
     BoxedClosure* getPassedClosure() { return passed_closure; }
@@ -584,9 +589,13 @@ Value ASTInterpreter::visit_langPrimitive(AST_LangPrimitive* node) {
         assert(node->args.size() == 1);
         assert(node->args[0]->type == AST_TYPE::Name);
 
-        RELEASE_ASSERT(source_info->ast->type == AST_TYPE::Module, "import * not supported in functions");
+        RELEASE_ASSERT(source_info->ast->type == AST_TYPE::Module || source_info->ast->type == AST_TYPE::Suite,
+                       "import * not supported in functions");
 
         Value module = visit_expr(node->args[0]);
+
+        RELEASE_ASSERT(globals == source_info->parent_module,
+                       "'import *' currently not supported with overridden globals");
         v = importStar(module.o, source_info->parent_module);
     } else if (node->opcode == AST_LangPrimitive::NONE) {
         v = None;
@@ -1314,6 +1323,12 @@ AST_stmt* getCurrentStatementForInterpretedFrame(void* frame_ptr) {
     return interpreter->getCurrentStatement();
 }
 
+Box* getGlobalsForInterpretedFrame(void* frame_ptr) {
+    ASTInterpreter* interpreter = s_interpreterMap[frame_ptr];
+    assert(interpreter);
+    return interpreter->getGlobals();
+}
+
 CompiledFunction* getCFForInterpretedFrame(void* frame_ptr) {
     ASTInterpreter* interpreter = s_interpreterMap[frame_ptr];
     assert(interpreter);

src/codegen/ast_interpreter.h

@@ -40,6 +40,7 @@ Box* astInterpretFrom(CompiledFunction* cf, AST_expr* after_expr, AST_stmt* encl
                       FrameStackState frame_state);
 
 AST_stmt* getCurrentStatementForInterpretedFrame(void* frame_ptr);
+Box* getGlobalsForInterpretedFrame(void* frame_ptr);
 CompiledFunction* getCFForInterpretedFrame(void* frame_ptr);
 struct FrameInfo;
 FrameInfo* getFrameInfoForInterpretedFrame(void* frame_ptr);

src/codegen/irgen/hooks.cpp

@@ -332,6 +332,8 @@ template <typename AST_Type>
 Box* evalOrExec(AST_Type* source, std::vector<AST_stmt*>& body, BoxedModule* bm, BoxedDict* globals, Box* boxedLocals) {
     CompiledFunction* cf;
 
+    assert(!globals || globals->cls == dict_cls);
+
     { // scope for limiting the locked region:
         LOCK_REGION(codegen_rwlock.asWrite());
 
@@ -361,6 +363,10 @@ Box* evalOrExec(AST_Type* source, std::vector<AST_stmt*>& body, BoxedModule* bm,
 Box* eval(Box* boxedCode) {
     Box* boxedLocals = fastLocalsToBoxedLocals();
     BoxedModule* module = getCurrentModule();
+    Box* globals = getGlobals();
+
+    if (globals == module)
+        globals = NULL;
 
     // TODO error message if parse fails or if it isn't an expr
     // TODO should have a cleaner interface that can parse the Expression directly
@@ -378,7 +384,9 @@ Box* eval(Box* boxedCode) {
     stmt->value = parsedExpr->body;
     std::vector<AST_stmt*> body = { stmt };
 
-    return evalOrExec<AST_Expression>(parsedExpr, body, module, NULL, boxedLocals);
+    assert(!globals || globals->cls == dict_cls);
+
+    return evalOrExec<AST_Expression>(parsedExpr, body, module, static_cast<BoxedDict*>(globals), boxedLocals);
 }
 
 Box* exec(Box* boxedCode, Box* globals, Box* locals) {
@@ -398,10 +406,22 @@ Box* exec(Box* boxedCode, Box* globals, Box* locals) {
         locals = fastLocalsToBoxedLocals();
     }
 
+    if (globals == NULL)
+        globals = getGlobals();
+
     BoxedModule* module = getCurrentModule();
 
+    if (globals == module)
+        globals = NULL;
+
     assert(!globals || globals->cls == dict_cls);
 
+    if (globals) {
+        // From CPython (they set it to be f->f_builtins):
+        if (PyDict_GetItemString(globals, "__builtins__") == NULL)
+            PyDict_SetItemString(globals, "__builtins__", builtins_module);
+    }
+
     // TODO same issues as in `eval`
     RELEASE_ASSERT(boxedCode->cls == str_cls, "");
     const char* code = static_cast<BoxedString*>(boxedCode)->s.c_str();

src/codegen/unwinding.cpp

@@ -300,6 +300,17 @@ public:
         abort();
     }
 
+    Box* getGlobals() {
+        if (id.type == PythonFrameId::COMPILED) {
+            CompiledFunction* cf = getCF();
+            assert(cf->clfunc->source->scoping->areGlobalsFromModule());
+            return cf->clfunc->source->parent_module;
+        } else if (id.type == PythonFrameId::INTERPRETED) {
+            return getGlobalsForInterpretedFrame((void*)id.bp);
+        }
+        abort();
+    }
+
     FrameInfo* getFrameInfo() {
         if (id.type == PythonFrameId::COMPILED) {
             CompiledFunction* cf = getCF();
@@ -561,6 +572,21 @@ CompiledFunction* getTopCompiledFunction() {
     return getTopPythonFrame()->getCF();
 }
 
+Box* getGlobals() {
+    auto it = getTopPythonFrame();
+    return it->getGlobals();
+}
+
+Box* getGlobalsDict() {
+    Box* globals = getGlobals();
+    if (!globals)
+        return NULL;
+
+    if (isSubclass(globals->cls, module_cls))
+        return makeAttrWrapper(globals);
+    return globals;
+}
+
 BoxedModule* getCurrentModule() {
     CompiledFunction* compiledFunction = getTopCompiledFunction();
     if (!compiledFunction)
@@ -656,7 +682,7 @@ Box* fastLocalsToBoxedLocals() {
         if (scope_info->areLocalsFromModule()) {
             // TODO we should cache this in frame_info->locals or something so that locals()
             // (and globals() too) will always return the same dict
-            return makeAttrWrapper(getCurrentModule());
+            return getGlobalsDict();
         }
 
         if (frame_iter.getId().type == PythonFrameId::COMPILED) {

src/codegen/unwinding.h

@@ -28,6 +28,8 @@ class BoxedTraceback;
 struct FrameInfo;
 
 BoxedModule* getCurrentModule();
+Box* getGlobals();     // returns either the module or a globals dict
+Box* getGlobalsDict(); // always returns a dict-like object
 
 BoxedTraceback* getTraceback();
 

src/core/cfg.cpp

@@ -2465,12 +2465,18 @@ CFG* computeCFG(SourceInfo* source, std::vector<AST_stmt*> body) {
 
     if (source->ast->type == AST_TYPE::ClassDef) {
         // A classdef always starts with "__module__ = __name__"
-        Box* module_name = source->parent_module->getattr("__name__", NULL);
-        assert(module_name->cls == str_cls);
         AST_Assign* module_assign = new AST_Assign();
         module_assign->targets.push_back(
             new AST_Name(source->getInternedStrings().get("__module__"), AST_TYPE::Store, source->ast->lineno));
-        module_assign->value = new AST_Str(static_cast<BoxedString*>(module_name)->s);
+
+        if (source->scoping->areGlobalsFromModule()) {
+            Box* module_name = source->parent_module->getattr("__name__", NULL);
+            assert(module_name->cls == str_cls);
+            module_assign->value = new AST_Str(static_cast<BoxedString*>(module_name)->s);
+        } else {
+            module_assign->value
+                = new AST_Name(source->getInternedStrings().get("__name__"), AST_TYPE::Load, source->ast->lineno);
+        }
         module_assign->lineno = 0;
         visitor.push_back(module_assign);
 

src/runtime/builtin_modules/builtins.cpp

@@ -725,9 +725,8 @@ public:
 };
 
 Box* globals() {
-    BoxedModule* m = getCurrentModule();
     // TODO is it ok that we don't return a real dict here?
-    return makeAttrWrapper(m);
+    return getGlobalsDict();
 }
 
 Box* locals() {
@@ -774,6 +773,7 @@ Box* execfile(Box* _fn) {
 
     // Run directly inside the current module:
     AST_Module* ast = caching_parse_file(fn->s.c_str());
+    ASSERT(getExecutionPoint().cf->clfunc->source->scoping->areGlobalsFromModule(), "need to pass custom globals in");
     compileAndRunModule(ast, getCurrentModule());
 
     return None;

src/runtime/import.cpp

@@ -442,6 +442,7 @@ 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(std::string* name, Box* globals, Box* from_imports, int level) {
+    assert(!globals || globals == None || isSubclass(globals->cls, module_cls));
     bool return_first = from_imports == None;
     assert(name);
 

src/runtime/objmodel.cpp

@@ -3917,8 +3917,12 @@ Box* typeNew(Box* _cls, Box* arg1, Box* arg2, Box** _args) {
         made->setattr(static_cast<BoxedString*>(k)->s, p.second, NULL);
     }
 
-    if (!made->hasattr("__module__"))
-        made->giveAttr("__module__", boxString(getCurrentModule()->name()));
+    if (!made->hasattr("__module__")) {
+        Box* gl = getGlobalsDict();
+        Box* attr = PyDict_GetItemString(gl, "__name__");
+        if (attr)
+            made->giveAttr("__module__", attr);
+    }
     if (!made->hasattr("__doc__"))
         made->giveAttr("__doc__", None);
 
@@ -4375,6 +4379,10 @@ extern "C" Box* importFrom(Box* _m, const std::string* name) {
 }
 
 extern "C" Box* importStar(Box* _from_module, BoxedModule* to_module) {
+    // TODO(kmod): it doesn't seem too bad to update this to take custom globals;
+    // it looks like mostly a matter of changing the getattr calls to getitem.
+    RELEASE_ASSERT(getGlobals() == to_module, "importStar doesn't support custom globals yet");
+
     assert(_from_module->cls == module_cls);
     BoxedModule* from_module = static_cast<BoxedModule*>(_from_module);
 

src/runtime/types.cpp

@@ -23,6 +23,7 @@
 
 #include "capi/typeobject.h"
 #include "capi/types.h"
+#include "codegen/unwinding.h"
 #include "core/options.h"
 #include "core/stats.h"
 #include "core/types.h"
@@ -273,7 +274,7 @@ extern "C" BoxedFunctionBase::BoxedFunctionBase(CLFunction* f)
     : in_weakreflist(NULL), f(f), closure(NULL), isGenerator(false), ndefaults(0), defaults(NULL), modname(NULL),
       name(NULL), doc(NULL) {
     if (f->source) {
-        this->modname = f->source->parent_module->getattr("__name__", NULL);
+        this->modname = PyDict_GetItemString(getGlobalsDict(), "__name__");
         this->doc = f->source->getDocString();
     } else {
         this->modname = boxStringPtr(&builtinStr);
@@ -296,7 +297,7 @@ extern "C" BoxedFunctionBase::BoxedFunctionBase(CLFunction* f, std::initializer_
     }
 
     if (f->source) {
-        this->modname = f->source->parent_module->getattr("__name__", NULL);
+        this->modname = PyDict_GetItemString(getGlobalsDict(), "__name__");
         this->doc = f->source->getDocString();
     } else {
         this->modname = boxStringPtr(&builtinStr);
@@ -676,8 +677,8 @@ extern "C" Box* createUserClass(const std::string* name, Box* _bases, Box* _attr
         // an error, then look up ob_type (aka cls)
         metaclass = bases->elts[0]->cls;
     } else {
-        BoxedModule* m = getCurrentModule();
-        metaclass = m->getattr("__metaclass__");
+        Box* gl = getGlobalsDict();
+        metaclass = PyDict_GetItemString(gl, "__metaclass__");
 
         if (!metaclass) {
             metaclass = classobj_cls;
@@ -978,7 +979,8 @@ Box* typeRepr(BoxedClass* self) {
     Box* m = self->getattr("__module__");
     if (m && m->cls == str_cls) {
         BoxedString* sm = static_cast<BoxedString*>(m);
-        os << sm->s << '.';
+        if (sm->s != "__builtin__")
+            os << sm->s << '.';
     }
 
     os << self->tp_name;

test/tests/exec_import_star_in.py

+# expected: fail
+
+g = {}
+exec """from import_target import *""" in g
+del g['__builtins__']
+print sorted(g.keys())

test/tests/exec_in_test.py

@@ -46,6 +46,7 @@ f(-1)
 for i in xrange(1000):
     g['b'] = 6
     l['f'](i)
+print l['f'].__module__
 
 print 'Test global access in comprehensions'
 
@@ -74,3 +75,46 @@ exec_("""global a
 print a
 print b""", g, l)
 
+
+exec """print __name__"""
+exec """print __name__""" in {}, {}
+
+
+# Test classdefs in execs:
+b = 3
+a = 2
+s = """class C(object):
+    print "b =", b
+    if b:
+        c = 2
+    else:
+        a = -1
+    print a, b
+print C.__module__, C.__name__, repr(C)
+"""
+exec s in {'a': 1, 'b': 5}, {'b': 2}
+exec s
+
+
+# Test old-style classdefs in execs:
+b = 3
+a = 2
+s = """class C():
+    print "b =", b
+    if b:
+        c = 2
+    else:
+        a = -1
+    print a, b
+print C.__module__
+"""
+exec s in {'a': 1, 'b': 5}, {'b': 2}
+exec s
+
+
+
+
+# test eval+exec in exec:
+a = 5
+exec """print eval('a')""" in {'a': 6}, {}
+exec """exec 'print a' """ in {'a': 6}, {}