diff --git a/Lib/test/test_freeze/test_module_proxy.py b/Lib/test/test_freeze/test_module_proxy.py index d94ccd538e3c7a..a3f7fb024e8ef6 100644 --- a/Lib/test/test_freeze/test_module_proxy.py +++ b/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() diff --git a/Modules/_randommodule.c b/Modules/_randommodule.c index aa2fd28c232f28..1d26f269bd3abe 100644 --- a/Modules/_randommodule.c +++ b/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; } diff --git a/Objects/moduleobject.c b/Objects/moduleobject.c index fb94c6a648f48b..3c6b76a0f0f7d9 100644 --- a/Objects/moduleobject.c +++ b/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 = { diff --git a/Python/immutability.c b/Python/immutability.c index 0c4a83716b3d3f..c4feb45d0511c7 100644 --- a/Python/immutability.c +++ b/Python/immutability.c @@ -174,6 +174,10 @@ int init_state(struct _Py_immutability_state *state) return -1; } } + + if (_PyImmutability_SetFreezable((PyObject*)&PyModule_Type, _Py_FREEZABLE_PROXY)) { + return -1; + } return 0; } @@ -1498,7 +1502,9 @@ int _PyImmutability_SetFreezable(PyObject *obj, _Py_freezable_status status) 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; diff --git a/Python/pystate.c b/Python/pystate.c index 5a10b241fd27b6..90ea2571941a56 100644 --- a/Python/pystate.c +++ b/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;