Adjusting _Py_FREEZABLE_PROXY

Lib/test/test_freeze/test_module_proxy.py

@@ -70,6 +70,20 @@ def coin():
             elif hasattr(mut_random, attr):
                 delattr(mut_random, attr)
 
+    def test_proxy_reimport(self):
+        import random
+        self.assertFalse(is_frozen(random))
+
+        freeze(random)
+
+        # Unimport "random"
+        old_mut_random = sys.mut_modules["random"]
+        del sys.mut_modules["random"]
+
+        # This should reimport the module
+        random.random()
+
+
 
 if __name__ == '__main__':
     unittest.main()

Modules/_randommodule.c

@@ -639,6 +639,13 @@ _random_exec(PyObject *module)
     if (state->Long___abs__ == NULL) {
         return -1;
     }
+
+    if (_PyImmutability_SetFreezable(module, _Py_FREEZABLE_NO) < 0
+        || _PyImmutability_SetFreezable(state->Random_Type, _Py_FREEZABLE_NO) < 0
+    ) {
+        return -1;
+    }
+
     return 0;
 }
 

Objects/moduleobject.c

@@ -1501,10 +1501,14 @@ static PyGetSetDef module_getsets[] = {
 static int
 module_reachable(PyObject *self, visitproc visit, void *arg)
 {
+    // FIXME(immutability): Allow modules to define their own custom
+    // `md_reachable` function. Currently, we're falling back on
+    // `md_traverse`
     Py_VISIT(_PyObject_CAST(Py_TYPE(self)));
     return module_traverse(self, visit, arg);
 }
 
+// Artifact[Implementation]: This turns an existing ModuleObject into a proxy object:
 static int
 module_make_immutable_proxy(PyObject *self) {
     // Use cast, since we want this exact object
@@ -1554,13 +1558,19 @@ module_make_immutable_proxy(PyObject *self) {
     return 0;
 }
 
-// Artifact[Implementation]: The pre-freeze hook of module objects
 static int
 module_prefreeze(PyObject *self) {
-    // TODO(immutable): Check if the module defines a custom pre-freeze hook:
+    // FIXME(immutability): Modules can define their own pre-freeze hook
+    // and then delegate to this function to turn themself into a
+    // proxy object. While this works, it's a bit cumbersome. There should
+    // be an easier and more direct way
 
-    // TODO(immutable): Check if the module wants to be a proxy first:
-    return module_make_immutable_proxy(self);
+    _Py_freezable_status status = _PyImmutability_GetFreezable(self);
+    if (status == _Py_FREEZABLE_PROXY) {
+        return module_make_immutable_proxy(self);
+    }
+
+    return 0;
 }
 
 PyTypeObject PyModule_Type = {

Python/immutability.c

@@ -174,6 +174,10 @@
             return -1;
         }
     }
+
+    if (_PyImmutability_SetFreezable((PyObject*)&PyModule_Type, _Py_FREEZABLE_PROXY)) {
+        return -1;
+    }
     return 0;
 }
 
@@ -671,7 +675,7 @@
     return true;
 }
 
 static PyObject* get_next(PyObject* obj, struct FreezeState *freeze_state)
 {
     (void)freeze_state;
     PyObject* next = scc_next(obj);
@@ -1498,7 +1502,9 @@
         return -1;
     }
 
-    if (status == _Py_FREEZABLE_PROXY && !PyModule_Check(obj)) {
+    if (status == _Py_FREEZABLE_PROXY
+        && !(PyModule_Check(obj) || obj == _PyObject_CAST(&PyModule_Type))
+    ) {
         PyErr_SetString(PyExc_TypeError,
                         "FREEZABLE_PROXY can only be set on module objects");
         return -1;
@@ -2124,7 +2130,7 @@
         return 0;
     }
 #else
 #error "Immutability currently only works on 64bit platforms"
 #endif
 
     // Mark pre-freeze hook as completed. This has to be set before calling

Python/pystate.c

@@ -1373,6 +1373,28 @@ PyModuleObject* _PyInterpreterState_GetModuleState(PyObject *mod) {
         PyModuleObject *self = (PyModuleObject*) mod;
 
         if (!PyDict_Contains(is->mutable_modules, self->md_name)) {
+            // Get the `sys.modules` reference
+            PyObject* modules = _PyImport_GetModulesRef(is);
+            if (modules == Py_None) {
+                return NULL;
+            }
+
+            PyObject *modules_mod = NULL;
+            if (PyDict_GetItemRef(modules, self->md_name, &modules_mod) < 0) {
+                return 0;
+            }
+            if (modules_mod == mod) {
+                // The modules in `sys.modules` is this frozen mod but there is
+                // no mutable state in `sys.mut_modules`, probably because it was
+                // unimported? Either way: 
+                // Remove `mod` from `sys.modules` and trigger a reimport.
+                if (PyDict_DelItem(modules, self->md_name) < 0) {
+                    Py_DECREF(modules_mod);
+                    return NULL;
+                }
+            }
+            Py_XDECREF(modules_mod);
+
             // Importing the module will import the module or return the already
             // imported instance in `sys.modules`.
             PyObject *local_mod = PyImport_Import(self->md_name);
@@ -1392,10 +1414,6 @@ PyModuleObject* _PyInterpreterState_GetModuleState(PyObject *mod) {
             }
 
             // Place immutable proxy in `sys.modules[dict]`
-            PyObject* modules = _PyImport_GetModulesRef(is);
-            if (modules == Py_None) {
-                return NULL;
-            }
             res = PyDict_SetItem(modules, self->md_name, _PyObject_CAST(self));
             if (res != 0) {
                 return NULL;