Working on global freezing

Signed-off-by: Matthew A Johnson <matjoh@microsoft.com>

Author: matajoh

Include/internal/pycore_global_objects_fini_generated.h

@@ -48,10 +48,12 @@ struct _Py_global_strings {
         STRUCT_FOR_STR(json_decoder, "json.decoder")
         STRUCT_FOR_STR(kwdefaults, ".kwdefaults")
         STRUCT_FOR_STR(list_err, "list index out of range")
+        STRUCT_FOR_STR(namedtuple, "namedtuple")
         STRUCT_FOR_STR(newline, "\n")
         STRUCT_FOR_STR(open_br, "{")
         STRUCT_FOR_STR(percent, "%")
         STRUCT_FOR_STR(shim_name, "<shim>")
+        STRUCT_FOR_STR(startswith, "startswith")
         STRUCT_FOR_STR(type_params, ".type_params")
         STRUCT_FOR_STR(utf_8, "utf-8")
     } literals;

Include/internal/pycore_global_strings.h

@@ -0,0 +1,10 @@
+import unittest
+
+
+class TestFreezeGlobals(unittest.TestCase):
+    def test_freeze_globals(self):
+        freezeglobals()
+
+if __name__ == "__main__":
+    unittest.main()
+    # Test the freezeglobals function

Include/internal/pycore_runtime_init_generated.h

@@ -1775,6 +1775,7 @@ long_to_decimal_string_internal(PyObject *aa,
     }
     size_a = _PyLong_DigitCount(a);
     negative = _PyLong_IsNegative(a);
+    kind = 0;
 
     /* quick and dirty pre-check for overflowing the decimal digit limit,
        based on the inequality 10/3 >= log2(10)
@@ -2006,6 +2007,7 @@ long_format_binary(PyObject *aa, int base, int alternate,
     }
     size_a = _PyLong_DigitCount(a);
     negative = _PyLong_IsNegative(a);
+    kind = 0;
 
     /* Compute a rough upper bound for the length of the string */
     switch (base) {

Lib/test/test_freezeglobals.py

@@ -601,8 +601,10 @@ void
 _PyModule_Clear(PyObject *m)
 {
     PyObject *d = ((PyModuleObject *)m)->md_dict;
-    if (d != NULL)
+    if (d != NULL){
+        _Py_ClearImmutable(d);
         _PyModule_ClearDict(d);
+    }
 }
 
 void

Objects/longobject.c

@@ -111,20 +111,27 @@ static PyObject* walk_function(PyObject* op, stack* frontier)
 
     _Py_SetImmutable(op);
 
-
     f = (PyFunctionObject*)op;
 
+
     globals = f->func_globals;
     builtins = f->func_builtins;
 
-    if(PyUnicode_CompareWithASCIIString(f->func_module, "_frozen_importlib") == 0){
-        // we don't want to freeze the importlib module
-        Py_RETURN_NONE;
-    }
-
     module = PyImport_Import(f->func_module);
     if(module == NULL){
-        return module;
+        // clear the exception so we can check if the module is a namedtuple
+        PyObject* exc = PyErr_GetRaisedException();
+        _Py_DECLARE_STR(namedtuple, "namedtuple");
+        _Py_DECLARE_STR(startswith, "startswith");
+        PyObject* res = PyObject_CallMethodOneArg(f->func_module, &_Py_STR(startswith), &_Py_STR(namedtuple));
+        if(Py_IsTrue(res)){
+            // namedtuple creates a fake module, which cannot be imported
+            Py_RETURN_NONE;
+        }
+
+        // not a namedtuple, so we need to set the exception
+        PyErr_SetRaisedException(exc);
+        return NULL;
     }
 
     if(PyModule_Check(module)){
@@ -294,18 +301,6 @@ static PyObject* walk_function(PyObject* op, stack* frontier)
 
 static int _freeze_visit(PyObject* obj, void* frontier)
 {
-    if(PyModule_Check(obj)){
-        const char* name = PyModule_GetName(obj);
-
-        if(strcmp(name, "_frozen_importlib_external") == 0){
-            return 0;
-        }
-
-        if(strcmp(name, "_frozen_importlib") == 0){
-            return 0;
-        }
-    }
-
     if(!_Py_IsImmutable(obj)){
         if(stack_push((stack*)frontier, obj)){
             PyErr_NoMemory();
@@ -332,9 +327,39 @@ PyObject* _Py_Freeze(PyObject* obj)
         return PyErr_NoMemory();
     }
 
+    PyObject* frozen_importlib = PyImport_ImportModule("_frozen_importlib");
+    if(frozen_importlib == NULL){
+        stack_free(frontier);
+        return NULL;
+    }
+
+    PyObject* blocking_on = PyObject_GetAttrString(frozen_importlib, "_blocking_on");
+    if(blocking_on == NULL){
+        Py_DECREF(frozen_importlib);
+        stack_free(frontier);
+        return NULL;
+    }
+
+    PyObject* module_locks = PyObject_GetAttrString(frozen_importlib, "_module_locks");
+    if(module_locks == NULL){
+        Py_DECREF(blocking_on);
+        Py_DECREF(frozen_importlib);
+        stack_free(frontier);
+        return NULL;
+    }
+
+    Py_DECREF(frozen_importlib);
+
     while(!stack_empty(frontier)){
         PyObject* item = stack_pop(frontier); // item.rc = x + 1
 
+        if(item == blocking_on ||
+           item == module_locks){
+            // the module lock and blocking on dictionaries must remain mutable or else
+            // we will not be able to import modules
+            continue;
+        }
+
         PyTypeObject* type = Py_TYPE(item);
         traverseproc traverse;
         PyObject* type_op = NULL;
@@ -354,6 +379,8 @@ PyObject* _Py_Freeze(PyObject* obj)
         if(PyFunction_Check(item)){
             PyObject* err = walk_function(item, frontier);
             if(!Py_IsNone(err)){
+                Py_DECREF(blocking_on);
+                Py_DECREF(module_locks);
                 stack_free(frontier);
                 return err;
             }
@@ -363,6 +390,8 @@ PyObject* _Py_Freeze(PyObject* obj)
         traverse = type->tp_traverse;
         if(traverse != NULL){
             if(traverse(item, (visitproc)_freeze_visit, frontier)){
+                Py_DECREF(blocking_on);
+                Py_DECREF(module_locks);
                 stack_free(frontier);
                 return NULL;
             }
@@ -373,12 +402,16 @@ PyObject* _Py_Freeze(PyObject* obj)
         if (!_Py_IsImmutable(type_op)){
             if (stack_push(frontier, type_op))
             {
+                Py_DECREF(blocking_on);
+                Py_DECREF(module_locks);
                 stack_free(frontier);
                 return PyErr_NoMemory();
             }
         }
     }
 
+    Py_DECREF(blocking_on);
+    Py_DECREF(module_locks);
     stack_free(frontier);
 
     Py_RETURN_NONE;
@@ -390,22 +423,12 @@ PyObject* _Py_FreezeGlobals()
     PyObject* main_dict;
     PyInterpreterState* main;
 
-    // we will need to exclude importlib from this process
-    PyObject* importlib = PyUnicode_FromString("importlib");
-    if(importlib == NULL){
-        return NULL;
-    }
-
     // go module by module and freeze their global dictionaries
     PyObject* modules = PyImport_GetModuleDict();
     Py_ssize_t size = PyDict_Size(modules);
     PyObject* keys = PyDict_Keys(modules);
     for(Py_ssize_t i = 0; i < size; i++){
         PyObject* key = PyList_GetItem(keys, i);
-        if (PyUnicode_Contains(key, importlib)){
-            continue;
-        }
-
         PyObject* module = PyDict_GetItem(modules, key);
         PyObject* globals = PyModule_GetDict(module);
         ret = _Py_Freeze(globals);
@@ -416,7 +439,6 @@ PyObject* _Py_FreezeGlobals()
     }
 
     Py_DECREF(keys);
-    Py_DECREF(importlib);
 
     main = PyInterpreterState_Main();
     if(main == NULL){

Objects/moduleobject.c

@@ -1340,6 +1340,7 @@ finalize_modules_delete_special(PyThreadState *tstate, int verbose)
     };
 
     PyInterpreterState *interp = tstate->interp;
+    _Py_ClearImmutable(interp->builtins);
     if (verbose) {
         PySys_WriteStderr("# clear builtins._\n");
     }

Python/immutability.c

@@ -842,6 +842,9 @@ void
 PyErr_PrintEx(int set_sys_last_vars)
 {
     PyThreadState *tstate = _PyThreadState_GET();
+    if(set_sys_last_vars){
+        _Py_ClearImmutable(tstate->interp->sysdict);
+    }
     _PyErr_PrintEx(tstate, set_sys_last_vars);
 }