return globals() as locals() for module-level scopes

src/analysis/scoping_analysis.cpp

@@ -105,6 +105,8 @@ public:
 
     bool isPassedToViaClosure(InternedString name) override { return false; }
 
+    bool areLocalsFromModule() override { return true; }
+
     InternedString mangleName(InternedString id) override { return id; }
     InternedString internString(llvm::StringRef s) override { abort(); }
 };
@@ -245,6 +247,8 @@ public:
         return usage->got_from_closure.count(name) > 0 || usage->passthrough_accesses.count(name) > 0;
     }
 
+    bool areLocalsFromModule() override { return false; }
+
     InternedString mangleName(const InternedString id) override {
         return pyston::mangleName(id, usage->private_name, usage->scoping->getInternedStrings());
     }

src/analysis/scoping_analysis.h

@@ -114,6 +114,8 @@ public:
     // `exec` or `eval` scope.
     virtual bool usesNameLookup() = 0;
 
+    virtual bool areLocalsFromModule() = 0;
+
     virtual InternedString mangleName(InternedString id) = 0;
     virtual InternedString internString(llvm::StringRef) = 0;
 };

src/codegen/ast_interpreter.cpp

@@ -1095,15 +1095,7 @@ Value ASTInterpreter::visit_name(AST_Name* node) {
             return Value();
         }
         case ScopeInfo::VarScopeType::NAME: {
-            assert(frame_info.boxedLocals->cls == dict_cls);
-            auto& d = static_cast<BoxedDict*>(frame_info.boxedLocals)->d;
-            auto it = d.find(boxString(node->id.str()));
-            if (it != d.end()) {
-                Box* value = it->second;
-                return value;
-            }
-
-            return getGlobal(source_info->parent_module, &node->id.str());
+            return boxedLocalsGet(frame_info.boxedLocals, node->id.c_str(), source_info->parent_module);
         }
         default:
             abort();
@@ -1170,7 +1162,6 @@ Box* astInterpretFunctionEval(CompiledFunction* cf, Box* boxedLocals) {
 
     ASTInterpreter interpreter(cf);
     interpreter.initArguments(0, NULL, NULL, NULL, NULL, NULL, NULL);
-    RELEASE_ASSERT(boxedLocals->cls == dict_cls, "we don't support non-dicts here yet");
     interpreter.setBoxedLocals(boxedLocals);
     Value v = ASTInterpreter::execute(interpreter);
 

src/codegen/unwinding.cpp

@@ -638,12 +638,20 @@ Box* fastLocalsToBoxedLocals() {
     for (PythonFrameIterator& frame_iter : unwindPythonFrames()) {
         BoxedDict* d;
         BoxedClosure* closure;
-        CompiledFunction* cf;
         FrameInfo* frame_info;
+
+        CompiledFunction* cf = frame_iter.getCF();
+        ScopeInfo* scope_info = cf->clfunc->source->getScopeInfo();
+
+        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());
+        }
+
         if (frame_iter.getId().type == PythonFrameId::COMPILED) {
             d = new BoxedDict();
 
-            cf = frame_iter.getCF();
             uint64_t ip = frame_iter.getId().ip;
 
             assert(ip > cf->code_start);
@@ -727,7 +735,6 @@ Box* fastLocalsToBoxedLocals() {
         } else if (frame_iter.getId().type == PythonFrameId::INTERPRETED) {
             d = localsForInterpretedFrame((void*)frame_iter.getId().bp, true);
             closure = passedClosureForInterpretedFrame((void*)frame_iter.getId().bp);
-            cf = getCFForInterpretedFrame((void*)frame_iter.getId().bp);
             frame_info = getFrameInfoForInterpretedFrame((void*)frame_iter.getId().bp);
         } else {
             abort();
@@ -747,7 +754,6 @@ Box* fastLocalsToBoxedLocals() {
                 const std::string& name = attr_offset.first();
                 int offset = attr_offset.second;
                 Box* val = closure->attrs.attr_list->attrs[offset];
-                ScopeInfo* scope_info = cf->clfunc->source->getScopeInfo();
                 if (val != NULL && scope_info->isPassedToViaClosure(scope_info->internString(name))) {
                     Box* boxedName = boxString(name);
                     if (d->d.count(boxedName) == 0) {

src/runtime/objmodel.cpp

@@ -4450,20 +4450,34 @@ Box* coerceUnicodeToStr(Box* unicode) {
     return r;
 }
 
+// TODO Make these fast, do inline caches and stuff
+
 extern "C" void boxedLocalsSet(Box* boxedLocals, const char* attr, Box* val) {
     setitem(boxedLocals, boxString(attr), val);
 }
 
 extern "C" Box* boxedLocalsGet(Box* boxedLocals, const char* attr, BoxedModule* parent_module) {
     assert(parent_module->cls == module_cls);
-
     assert(boxedLocals != NULL);
-    RELEASE_ASSERT(boxedLocals->cls == dict_cls, "we don't support non-dict here yet");
-    auto& d = static_cast<BoxedDict*>(boxedLocals)->d;
-    auto it = d.find(boxString(attr));
-    if (it != d.end()) {
-        Box* value = it->second;
-        return value;
+
+    if (boxedLocals->cls == dict_cls) {
+        auto& d = static_cast<BoxedDict*>(boxedLocals)->d;
+        auto it = d.find(boxString(attr));
+        if (it != d.end()) {
+            Box* value = it->second;
+            return value;
+        }
+    } else {
+        try {
+            return getitem(boxedLocals, boxString(attr));
+        } catch (ExcInfo e) {
+            // TODO should check the exact semantic here but it's something like:
+            // If it throws a KeyError, then the variable doesn't exist so move on
+            // and check the globals (below); otherwise, just propogate the exception.
+            if (!isInstance(e.value, KeyError)) {
+                throw;
+            }
+        }
     }
 
     // TODO exception name?

test/tests/eval_test.py

@@ -5,17 +5,17 @@ print eval("3 + 4")
 a = 5
 print eval("a")
 
-#print eval("[b for b in range(5)]")
-#print b
+print eval("[b for b in range(5)]")
+print b
 
-#c = 2
-#print eval("[c for c in range(5)]")
-#print c
+c = 2
+print eval("[c for c in range(5)]")
+print c
 
-#try:
-#    print eval("int('abc')")
-#except ValueError:
-#    print 'got ValueError'
+try:
+    print eval("int('abc')")
+except ValueError:
+    print 'got ValueError'
 
 d = 19
 e = 20
@@ -85,7 +85,7 @@ o = 300
 print 'eval eval o', eval("eval('o')")
 
 # This works in the global scope but not in the local scope, because o1 is a global:
-#print eval("[(lambda p1 : p1 + o1)(5) for o1 in range(5)]")
+# print eval("[(lambda p1 : p1 + o1)(5) for o1 in range(5)]")
 def lambda_func():
     try:
         pass #print eval("[(lambda p2 : p2 + o2)(5) for o2 in range(5)]")

test/tests/globals_func.py

@@ -18,3 +18,8 @@ except NameError:
 # You're allowed to assign through globals and have it affect the module:
 globals()['x'] = 1
 print x
+
+# locals should do the same as globals
+print locals()['x']
+locals()['x'] = 2
+print x