Fall back to __main__.__dict__ if running main.

Author: ericsnowcurrently

Include/internal/pycore_ceval.h

@@ -239,6 +239,7 @@ static inline void _Py_LeaveRecursiveCall(void)  {
 
 extern _PyInterpreterFrame* _PyEval_GetFrame(void);
 
+extern PyObject * _PyEval_GetGlobalsFromRunningMain(PyThreadState *);
 extern int _PyEval_EnsureBuiltins(
     PyThreadState *,
     PyObject *,

Objects/object.c

@@ -2083,9 +2083,23 @@ _dir_locals(void)
     PyObject *names;
     PyObject *locals;
 
-    locals = _PyEval_GetFrameLocals();
-    if (locals == NULL)
+    if (_PyEval_GetFrame() != NULL) {
+        locals = _PyEval_GetFrameLocals();
+    }
+    PyThreadState *tstate = _PyThreadState_GET();
+    locals = _PyEval_GetGlobalsFromRunningMain(tstate);
+    if (locals == NULL) {
+        if (!_PyErr_Occurred(tstate)) {
+            locals = _PyEval_GetFrameLocals();
+            assert(_PyErr_Occurred(tstate));
+        }
+    }
+    else {
+        Py_INCREF(locals);
+    }
+    if (locals == NULL) {
         return NULL;
+    }
 
     names = PyMapping_Keys(locals);
     Py_DECREF(locals);

Python/bltinmodule.c

@@ -983,10 +983,16 @@ builtin_eval_impl(PyObject *module, PyObject *source, PyObject *globals,
         Py_INCREF(globals);
     }
     else {
-        PyErr_SetString(PyExc_TypeError,
-            "eval must be given globals and locals "
-            "when called without a frame");
-        return NULL;
+        globals = _PyEval_GetGlobalsFromRunningMain(tstate);
+        if (globals == NULL) {
+            if (!_PyErr_Occurred(tstate)) {
+                PyErr_SetString(PyExc_TypeError,
+                    "eval must be given globals and locals "
+                    "when called without a frame");
+            }
+            return NULL;
+        }
+        Py_INCREF(globals);
     }
 
     if (locals != Py_None) {
@@ -1090,9 +1096,15 @@ builtin_exec_impl(PyObject *module, PyObject *source, PyObject *globals,
         Py_INCREF(globals);
     }
     else {
-        PyErr_SetString(PyExc_SystemError,
-                        "globals and locals cannot be NULL");
-        goto error;
+        globals = _PyEval_GetGlobalsFromRunningMain(tstate);
+        if (globals == NULL) {
+            if (!_PyErr_Occurred(tstate)) {
+                PyErr_SetString(PyExc_SystemError,
+                                "globals and locals cannot be NULL");
+            }
+            goto error;
+        }
+        Py_INCREF(globals);
     }
 
     if (locals != Py_None) {
@@ -1255,12 +1267,21 @@ static PyObject *
 builtin_globals_impl(PyObject *module)
 /*[clinic end generated code: output=e5dd1527067b94d2 input=9327576f92bb48ba]*/
 {
+    PyObject *globals;
     if (_PyEval_GetFrame() != NULL) {
-        PyObject *globals = PyEval_GetGlobals();
+        globals = PyEval_GetGlobals();
         assert(globals != NULL);
         return Py_NewRef(globals);
     }
-    Py_RETURN_NONE;
+    PyThreadState *tstate = _PyThreadState_GET();
+    globals = _PyEval_GetGlobalsFromRunningMain(tstate);
+    if (globals == NULL) {
+        if (_PyErr_Occurred(tstate)) {
+            return NULL;
+        }
+        Py_RETURN_NONE;
+    }
+    return Py_NewRef(globals);
 }
 
 
@@ -1910,7 +1931,15 @@ builtin_locals_impl(PyObject *module)
         assert(locals != NULL || PyErr_Occurred());
         return locals;
     }
-    Py_RETURN_NONE;
+    PyThreadState *tstate = _PyThreadState_GET();
+    locals = _PyEval_GetGlobalsFromRunningMain(tstate);
+    if (locals == NULL) {
+        if (_PyErr_Occurred(tstate)) {
+            return NULL;
+        }
+        Py_RETURN_NONE;
+    }
+    return Py_NewRef(locals);
 }
 
 
@@ -2646,7 +2675,20 @@ builtin_vars(PyObject *self, PyObject *args)
     if (!PyArg_UnpackTuple(args, "vars", 0, 1, &v))
         return NULL;
     if (v == NULL) {
-        d = _PyEval_GetFrameLocals();
+        if (_PyEval_GetFrame() != NULL) {
+            d = _PyEval_GetFrameLocals();
+        }
+        PyThreadState *tstate = _PyThreadState_GET();
+        d = _PyEval_GetGlobalsFromRunningMain(tstate);
+        if (d == NULL) {
+            if (!_PyErr_Occurred(tstate)) {
+                d = _PyEval_GetFrameLocals();
+                assert(_PyErr_Occurred(tstate));
+            }
+        }
+        else {
+            Py_INCREF(d);
+        }
     }
     else {
         if (PyObject_GetOptionalAttr(v, &_Py_ID(__dict__), &d) == 0) {

Python/ceval.c

@@ -2763,6 +2763,22 @@ PyEval_GetGlobals(void)
     return _PyEval_GetGlobals(tstate);
 }
 
+PyObject *
+_PyEval_GetGlobalsFromRunningMain(PyThreadState *tstate)
+{
+    if (!_PyInterpreterState_IsRunningMain(tstate->interp)) {
+        return NULL;
+    }
+    PyObject *mod = _Py_GetMainModule(tstate);
+    if (_Py_CheckMainModule(mod) < 0) {
+        Py_XDECREF(mod);
+        return NULL;
+    }
+    PyObject *globals = PyModule_GetDict(mod);  // borrowed
+    Py_DECREF(mod);
+    return globals;
+}
+
 int
 _PyEval_EnsureBuiltins(PyThreadState *tstate, PyObject *globals, int usemod,
                        PyObject **p_builtins)