From 9170dd3da445435362ac9d3bb58c7dc78799ac17 Mon Sep 17 00:00:00 2001 From: xFrednet Date: Thu, 5 Mar 2026 17:31:25 +0100 Subject: [PATCH 1/3] Immutability: Clear collecting flag during freezing --- Include/internal/pycore_gc.h | 10 ++++++++++ Python/immutability.c | 12 +++++++++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/Include/internal/pycore_gc.h b/Include/internal/pycore_gc.h index a6519aa086309d..2dfce32237a83c 100644 --- a/Include/internal/pycore_gc.h +++ b/Include/internal/pycore_gc.h @@ -205,6 +205,16 @@ static inline void _PyGC_CLEAR_FINALIZED(PyObject *op) { #endif } +static inline void _PyGC_CLEAR_COLLECTING(PyObject *op) { +#ifdef Py_GIL_DISABLED + // TODO(immutable): Does NoGil have a collecting flag? If so, how do we + // clear it? +#else + PyGC_Head *gc = _Py_AS_GC(op); + gc->_gc_prev &= ~_PyGC_PREV_MASK_COLLECTING; +#endif +} + /* Tell the GC to track this object. * diff --git a/Python/immutability.c b/Python/immutability.c index 6fd06a312ab0be..597d0e13dce7ef 100644 --- a/Python/immutability.c +++ b/Python/immutability.c @@ -512,6 +512,16 @@ static void scc_init(PyObject* obj) _PyObject_GC_UNTRACK(obj); } + // The GC uses the collecting flag to identify objects part of the + // current collection set. This flag remains while the finalizer + // of unreachable objects is being called. + // + // If something calls `freeze(obj)` as part of their finalizer we + // might receive an object with the flag set. This removes the flag + // to prevent future GC collections to assume this object is currently + // being collected. + _PyGC_CLEAR_COLLECTING(obj); + // Mark as pending so we can detect back edges in the traversal. IMMUTABLE_FLAG_FIELD(obj) |= _Py_IMMUTABLE_PENDING; @@ -678,7 +688,7 @@ static void scc_reset_root_refcount(PyObject* obj) } // This will restore the reference counts for the interior edges of the SCC. -// It calculates some properites of the SCC, to decide how it might be +// It calculates some properties of the SCC, to decide how it might be // finalised. Adds an RC to every element in the SCC. static void scc_add_internal_refcounts(PyObject* obj, struct SCCDetails* details) { From 029aad7b9c58b50359f6f50f21e25c6cfc9ed144 Mon Sep 17 00:00:00 2001 From: xFrednet Date: Thu, 5 Mar 2026 17:56:27 +0100 Subject: [PATCH 2/3] Immutability: Basic Pre-freeze infrastructure --- Include/cpython/object.h | 3 + Include/object.h | 1 + Lib/test/test_freeze/test_prefreeze.py | 58 ++++++++ Lib/test/test_sys.py | 2 +- Objects/funcobject.c | 110 ++++++++++++++ Objects/moduleobject.c | 104 +++++++------ Objects/weakrefobject.c | 2 + Python/immutability.c | 195 ++++++++----------------- 8 files changed, 290 insertions(+), 185 deletions(-) create mode 100644 Lib/test/test_freeze/test_prefreeze.py diff --git a/Include/cpython/object.h b/Include/cpython/object.h index 7f62f79fb421a4..147e48a602179e 100644 --- a/Include/cpython/object.h +++ b/Include/cpython/object.h @@ -242,6 +242,9 @@ struct _typeobject { /* call function for all referenced objects (includes non-cyclic refs) */ traverseproc tp_reachable; + + /* A callback called before a type is frozen. */ + prefreezeproc tp_prefreeze; }; #define _Py_ATTR_CACHE_UNUSED (30000) // (see tp_versions_used) diff --git a/Include/object.h b/Include/object.h index 9585f4a1d67a52..5e5929e19eeb39 100644 --- a/Include/object.h +++ b/Include/object.h @@ -358,6 +358,7 @@ typedef int(*objobjargproc)(PyObject *, PyObject *, PyObject *); typedef int (*objobjproc)(PyObject *, PyObject *); typedef int (*visitproc)(PyObject *, void *); typedef int (*traverseproc)(PyObject *, visitproc, void *); +typedef int (*prefreezeproc)(PyObject *); typedef void (*freefunc)(void *); diff --git a/Lib/test/test_freeze/test_prefreeze.py b/Lib/test/test_freeze/test_prefreeze.py new file mode 100644 index 00000000000000..c2ffad26f97d41 --- /dev/null +++ b/Lib/test/test_freeze/test_prefreeze.py @@ -0,0 +1,58 @@ +import unittest + +from immutable import NotFreezable, freeze, isfrozen + + +class TestPreFreezeHook(unittest.TestCase): + def test_prefreeze_hook_is_called(self): + class C: + def __init__(self): + self.hook_calls = 0 + + def __pre_freeze__(self): + self.hook_calls += 1 + + obj = C() + freeze(obj) + + self.assertEqual(obj.hook_calls, 1) + self.assertTrue(isfrozen(obj)) + + def test_prefreeze_hook_runs_before_object_is_frozen(self): + class C: + def __init__(self): + self.was_frozen_inside_hook = None + + def __pre_freeze__(self): + self.was_frozen_inside_hook = isfrozen(obj) + + obj = C() + freeze(obj) + + self.assertIs(obj.was_frozen_inside_hook, False) + self.assertTrue(isfrozen(obj)) + + def test_prefreeze_hook_remains_called_after_failure(self): + class C: + def __init__(self): + self.hook_calls = 0 + self.child = NotFreezable() + + def __pre_freeze__(self): + self.hook_calls += 1 + + obj = C() + + with self.assertRaises(TypeError): + freeze(obj) + with self.assertRaises(TypeError): + freeze(obj) + with self.assertRaises(TypeError): + freeze(obj) + + self.assertEqual(obj.hook_calls, 1) + self.assertFalse(isfrozen(obj)) + + +if __name__ == "__main__": + unittest.main() diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py index 8cf4d062a673dd..f14d20acc6b168 100644 --- a/Lib/test/test_sys.py +++ b/Lib/test/test_sys.py @@ -1782,7 +1782,7 @@ def delx(self): del self.__x check((1,2,3), vsize('') + self.P + 3*self.P) # type # static type: PyTypeObject - fmt = 'P2nPI13Pl4Pn9Pn12PIPcP' + fmt = 'P2nPI13Pl4Pn9Pn12PIPcPP' s = vsize(fmt) check(int, s) typeid = 'n' if support.Py_GIL_DISABLED else '' diff --git a/Objects/funcobject.c b/Objects/funcobject.c index 7a8b1c91d60e52..e728b17e96baca 100644 --- a/Objects/funcobject.c +++ b/Objects/funcobject.c @@ -1212,6 +1212,115 @@ func_reachable(PyObject *self, visitproc visit, void *arg) return func_traverse(self, visit, arg); } +/** + * Special function for replacing globals and builtins with a copy of just what they use. + * + * This is necessary because the function object has a pointer to the global + * dictionary, and this is problematic because freezing any function directly + * (as we do with other objects) would make all globals immutable. + * + * Instead, we walk the function and find any places where it references + * global variables or builtins, and then freeze just those objects. The globals + * and builtins dictionaries for the function are then replaced with + * copies containing just those globals and builtins we were able to determine + * the function uses. + */ +static int func_prefreeze_shadow_captures(PyObject* op) +{ + _PyObject_ASSERT(op, PyFunction_Check(op)); + + PyObject* shadow_builtins = NULL; + PyObject* shadow_globals = NULL; + PyObject* shadow_closure = NULL; + Py_ssize_t size; + + PyFunctionObject* f = _PyFunction_CAST(op); + PyObject* globals = f->func_globals; + PyObject* builtins = f->func_builtins; + + shadow_builtins = PyDict_New(); + if(shadow_builtins == NULL){ + goto nomemory; + } + + shadow_globals = PyDict_New(); + if(shadow_globals == NULL){ + goto nomemory; + } + + if (PyDict_SetItemString(shadow_globals, "__builtins__", Py_NewRef(shadow_builtins))) { + goto error; + } + + _PyObject_ASSERT(f->func_code, PyCode_Check(f->func_code)); + PyCodeObject* f_code = (PyCodeObject*)f->func_code; + + size = 0; + if (f_code->co_names != NULL) + size = PySequence_Fast_GET_SIZE(f_code->co_names); + for (Py_ssize_t i = 0; i < size; i++) { + PyObject* name = PySequence_Fast_GET_ITEM(f_code->co_names, i); + + if( PyDict_Contains(globals, name)){ + PyObject* value = PyDict_GetItem(globals, name); + if (PyDict_SetItem(shadow_globals, Py_NewRef(name), Py_NewRef(value))) { + goto error; + } + } else if (PyDict_Contains(builtins, name)) { + PyObject* value = PyDict_GetItem(builtins, name); + if (PyDict_SetItem(shadow_builtins, Py_NewRef(name), Py_NewRef(value))) { + goto error; + } + } + } + + // Shadow cells with a new frozen cell to warn on reassignments in the + // capturing function. + size = 0; + if(f->func_closure != NULL) { + size = PyTuple_Size(f->func_closure); + if (size == -1) { + goto error; + } + shadow_closure = PyTuple_New(size); + if (shadow_closure == NULL) { + goto error; + } + } + for(Py_ssize_t i=0; i < size; ++i){ + PyObject* cellvar = PyTuple_GET_ITEM(f->func_closure, i); + PyObject* value = PyCell_GET(cellvar); + + PyObject* shadow_cellvar = PyCell_New(value); + if(PyTuple_SetItem(shadow_closure, i, shadow_cellvar) == -1){ + goto error; + } + } + + if (f->func_annotations == NULL) { + PyObject* new_annotations = PyDict_New(); + if (new_annotations == NULL) { + goto nomemory; + } + f->func_annotations = new_annotations; + } + + // Only assign them at the end when everything succeeded + Py_XSETREF(f->func_closure, shadow_closure); + Py_SETREF(f->func_globals, shadow_globals); + Py_SETREF(f->func_builtins, shadow_builtins); + + return 0; + +nomemory: + PyErr_NoMemory(); +error: + Py_XDECREF(shadow_builtins); + Py_XDECREF(shadow_globals); + Py_XDECREF(shadow_closure); + return -1; +} + PyTypeObject PyFunction_Type = { PyVarObject_HEAD_INIT(&PyType_Type, 0) "function", @@ -1254,6 +1363,7 @@ PyTypeObject PyFunction_Type = { 0, /* tp_alloc */ func_new, /* tp_new */ .tp_reachable = func_reachable, + .tp_prefreeze = func_prefreeze_shadow_captures, }; diff --git a/Objects/moduleobject.c b/Objects/moduleobject.c index f97d7f66d90a14..9b97ce5c0e164f 100644 --- a/Objects/moduleobject.c +++ b/Objects/moduleobject.c @@ -864,54 +864,6 @@ _PyModule_ClearDict(PyObject *d) } -int _Py_module_freeze_hook(PyObject *self) { - // Use cast, since we want this exact object - PyModuleObject *m = _PyModule_CAST(self); - - if (m->md_frozen) { - return 0; - } - - // Get the interpreter state early to make error handling easy - PyInterpreterState* ip = PyInterpreterState_Get(); - if (ip == NULL) { - PyErr_Format(PyExc_RuntimeError, "Well, this is a problem", Py_None); - return -1; - } - - // Create a new module module - PyModuleObject *mut_state = new_module_notrack(&PyModule_Type); - if (mut_state == NULL) { - PyErr_NoMemory(); - return -1; - } - track_module(mut_state); - - // Insert our mutable module into `sys.mut_modules` - if (PyDict_SetItem(ip->mutable_modules, m->md_name, _PyObject_CAST(mut_state))) { - // Make sure failure keeps self intact - Py_DECREF(mut_state); - return -1; - } - - // Copy mutable state - mut_state->md_name = Py_NewRef(m->md_name); - mut_state->md_dict = m->md_dict; - mut_state->md_def = m->md_def; - mut_state->md_state = m->md_state; - mut_state->md_weaklist = m->md_weaklist; - - // Clear the state to freeze the module - m->md_dict = NULL; - m->md_def = NULL; - m->md_state = NULL; - m->md_weaklist = NULL; - m->md_frozen = true; - m->ob_base.ob_type = &_PyImmModule_Type; - - return 0; -} - /*[clinic input] class module "PyModuleObject *" "&PyModule_Type" [clinic start generated code]*/ @@ -1553,6 +1505,61 @@ module_reachable(PyObject *self, visitproc visit, void *arg) return module_traverse(self, visit, arg); } +int module_make_immutable_proxy(PyObject *self) { + // Use cast, since we want this exact object + PyModuleObject *m = _PyModule_CAST(self); + + if (m->md_frozen) { + return 0; + } + + // Get the interpreter state early to make error handling easy + PyInterpreterState* ip = PyInterpreterState_Get(); + if (ip == NULL) { + PyErr_Format(PyExc_RuntimeError, "Well, this is a problem", Py_None); + return -1; + } + + // Create a new module module + PyModuleObject *mut_state = new_module_notrack(&PyModule_Type); + if (mut_state == NULL) { + PyErr_NoMemory(); + return -1; + } + track_module(mut_state); + + // Insert our mutable module into `sys.mut_modules` + if (PyDict_SetItem(ip->mutable_modules, m->md_name, _PyObject_CAST(mut_state))) { + // Make sure failure keeps self intact + Py_DECREF(mut_state); + return -1; + } + + // Copy mutable state + mut_state->md_name = Py_NewRef(m->md_name); + mut_state->md_dict = m->md_dict; + mut_state->md_def = m->md_def; + mut_state->md_state = m->md_state; + mut_state->md_weaklist = m->md_weaklist; + + // Clear the state to freeze the module + m->md_dict = NULL; + m->md_def = NULL; + m->md_state = NULL; + m->md_weaklist = NULL; + m->md_frozen = true; + m->ob_base.ob_type = &_PyImmModule_Type; + + return 0; +} + +int module_prefreeze(PyObject *self) { + // TODO(immutable): Check if the module defines a custom pre-freeze hook: + + // TODO(immutable): Check if the module wants to be a proxy first: + return module_make_immutable_proxy(self); +} + PyTypeObject PyModule_Type = { PyVarObject_HEAD_INIT(&PyType_Type, 0) "module", /* tp_name */ @@ -1595,6 +1602,7 @@ PyTypeObject PyModule_Type = { new_module, /* tp_new */ PyObject_GC_Del, /* tp_free */ .tp_reachable = module_reachable, + .tp_prefreeze = module_prefreeze, }; PyTypeObject _PyImmModule_Type = { diff --git a/Objects/weakrefobject.c b/Objects/weakrefobject.c index f6af75951ba0b2..2bcaf30217bf50 100644 --- a/Objects/weakrefobject.c +++ b/Objects/weakrefobject.c @@ -158,6 +158,8 @@ static int weakref_reachable(PyObject *op, visitproc visit, void *arg) { Py_VISIT(_PyObject_CAST(Py_TYPE(op))); + // FIXME(immutable): Currently we manually visit the weak references object. + // we might want to do this here to remove (some) special handling. return gc_traverse(op, visit, arg); } diff --git a/Python/immutability.c b/Python/immutability.c index 597d0e13dce7ef..0da016a1f2008d 100644 --- a/Python/immutability.c +++ b/Python/immutability.c @@ -1013,122 +1013,6 @@ static void mark_all_frozen(struct FreezeState *state) #endif } -/** - * Special function for replacing globals and builtins with a copy of just what they use. - * - * This is necessary because the function object has a pointer to the global - * dictionary, and this is problematic because freezing any function directly - * (as we do with other objects) would make all globals immutable. - * - * Instead, we walk the function and find any places where it references - * global variables or builtins, and then freeze just those objects. The globals - * and builtins dictionaries for the function are then replaced with - * copies containing just those globals and builtins we were able to determine - * the function uses. - */ -static int shadow_function_globals(PyObject* op) -{ - _PyObject_ASSERT(op, PyFunction_Check(op)); - - PyObject* shadow_builtins = NULL; - PyObject* shadow_globals = NULL; - PyObject* shadow_closure = NULL; - Py_ssize_t size; - - PyFunctionObject* f = _PyFunction_CAST(op); - PyObject* globals = f->func_globals; - PyObject* builtins = f->func_builtins; - - shadow_builtins = PyDict_New(); - if(shadow_builtins == NULL){ - goto nomemory; - } - - debug("Shadowing builtins for function %s (%p)\n", f->func_name, f); - debug(" Original builtins: %p\n", builtins); - debug(" Shadow builtins: %p\n", shadow_builtins); - - shadow_globals = PyDict_New(); - if(shadow_globals == NULL){ - goto nomemory; - } - debug("Shadowing globals for function %s (%p)\n", f->func_name, f); - debug(" Original globals: %p\n", globals); - debug(" Shadow globals: %p\n", shadow_globals); - - if (PyDict_SetItemString(shadow_globals, "__builtins__", Py_NewRef(shadow_builtins))) { - goto error; - } - - _PyObject_ASSERT(f->func_code, PyCode_Check(f->func_code)); - PyCodeObject* f_code = (PyCodeObject*)f->func_code; - - size = 0; - if (f_code->co_names != NULL) - size = PySequence_Fast_GET_SIZE(f_code->co_names); - for (Py_ssize_t i = 0; i < size; i++) { - PyObject* name = PySequence_Fast_GET_ITEM(f_code->co_names, i); - - if( PyDict_Contains(globals, name)){ - PyObject* value = PyDict_GetItem(globals, name); - if (PyDict_SetItem(shadow_globals, Py_NewRef(name), Py_NewRef(value))) { - goto error; - } - } else if (PyDict_Contains(builtins, name)) { - PyObject* value = PyDict_GetItem(builtins, name); - if (PyDict_SetItem(shadow_builtins, Py_NewRef(name), Py_NewRef(value))) { - goto error; - } - } - } - - // Shadow cells with a new frozen cell to warn on reassignments in the - // capturing function. - size = 0; - if(f->func_closure != NULL) { - size = PyTuple_Size(f->func_closure); - if (size == -1) { - goto error; - } - shadow_closure = PyTuple_New(size); - if (shadow_closure == NULL) { - goto error; - } - } - for(Py_ssize_t i=0; i < size; ++i){ - PyObject* cellvar = PyTuple_GET_ITEM(f->func_closure, i); - PyObject* value = PyCell_GET(cellvar); - - PyObject* shadow_cellvar = PyCell_New(value); - if(PyTuple_SetItem(shadow_closure, i, shadow_cellvar) == -1){ - goto error; - } - } - - if (f->func_annotations == NULL) { - PyObject* new_annotations = PyDict_New(); - if (new_annotations == NULL) { - goto nomemory; - } - f->func_annotations = new_annotations; - } - - // Only assign them at the end when everything succeeded - Py_XSETREF(f->func_closure, shadow_closure); - Py_SETREF(f->func_globals, shadow_globals); - Py_SETREF(f->func_builtins, shadow_builtins); - - return 0; - -nomemory: - PyErr_NoMemory(); -error: - Py_XDECREF(shadow_builtins); - Py_XDECREF(shadow_globals); - Py_XDECREF(shadow_closure); - return -1; -} - static int freeze_visit(PyObject* obj, void* freeze_state_untyped) { struct FreezeState* freeze_state = (struct FreezeState *)freeze_state_untyped; @@ -1686,6 +1570,53 @@ void _Py_RefcntAdd_Immutable(PyObject *op, Py_ssize_t increment) // Macro that jumps to error, if the expression `x` does not succeed. #define SUCCEEDS(x) { do { int r = (x); if (r != 0) goto error; } while (0); } +static int run_pre_freeze_hook(PyObject* obj) { + PyObject *attr = NULL; + + // Skip Python-level hook lookup for type objects. For classes, + // `__pre_freeze__` resolves to an unbound function and calling it as a + // normal bound method would fail with a missing 'self' argument. + if (PyType_Check(obj)) { + return 0; + } + + // 0. Check if the pre-freeze hook already ran for this object + // TODO(immutable): Remember called hooks and return early + + // 1. Pre-freeze hooks are never called for shallow immutable objects + // TODO(immutable): Return early for shallow immutable. + + // 1. Check for the `__pre_freeze__` name + int res = PyObject_GetOptionalAttr(obj, &_Py_ID(__pre_freeze__), &attr); + if (res == -1) { + return -1; + } else if (res == 1) { + if (!PyCallable_Check(attr)) { + PyErr_Format( + PyExc_TypeError, + "'%.200s.__pre_freeze__' is not callable", + Py_TYPE(obj)->tp_name); + Py_DECREF(attr); + return -1; + } + PyObject *result = PyObject_CallNoArgs(attr); + Py_DECREF(attr); + if (result == NULL) { + return -1; + } + Py_DECREF(result); + } + + // 2. Check the type for `tp_prefreeze` + prefreezeproc prefreeze = Py_TYPE(obj)->tp_prefreeze; + if (prefreeze != NULL) { + return prefreeze(obj); + } + + // No pre-freeze hook, so we're good to go. + return 0; +} + static int traverse_freeze(PyObject* obj, struct FreezeState* freeze_state) { // WARNING @@ -1705,26 +1636,6 @@ static int traverse_freeze(PyObject* obj, struct FreezeState* freeze_state) return 0; } - PyObject *attr = NULL; - if (PyObject_GetOptionalAttr(obj, &_Py_ID(__pre_freeze__), &attr) == 1) - { - PyErr_SetString(PyExc_TypeError, "Pre-freeze hocks are currently WIP"); - Py_XDECREF(attr); - return -1; - } - Py_XDECREF(attr); - - // Function require some work to freeze, so we do not freeze the - // world as they mention globals and builtins. This will shadow what they - // use, and then we can freeze the those components. - if(PyFunction_Check(obj)){ - SUCCEEDS(shadow_function_globals(obj)); - } - - if (PyModule_Check(obj)) { - SUCCEEDS(_Py_module_freeze_hook(obj)); - } - SUCCEEDS(get_reachable_proc(Py_TYPE(obj))(obj, (visitproc)freeze_visit, freeze_state)); // Weak references are not followed by the GC, but should be @@ -1733,6 +1644,10 @@ static int traverse_freeze(PyObject* obj, struct FreezeState* freeze_state) if (PyWeakref_Check(obj)) { // Make the weak reference strong. // Get Ref increments the refcount. + // + // This could be done via a pre-freeze hook, but we only want to keep + // the strong reference if freezing succeeds. Having this as a special + // case makes this easier to handle. PyObject* wr; int res = PyWeakref_GetRef(obj, &wr); if (res == -1) { @@ -1880,6 +1795,14 @@ int _PyImmutability_Freeze(PyObject* obj) // New object, check if freezable SUCCEEDS(check_freezable(imm_state, item, &freeze_state)); + // Call the pre-freeze hook if one is present + SUCCEEDS(run_pre_freeze_hook(item)); + + // If the pre-freeze hook turned the object immutable, we want to skip it. + if (_Py_IsImmutable(item)) { + continue; + } + // Add to visited before putting in internal datastructures, so don't have // to account of internal RC manipulations. add_visited(item, &freeze_state); From ce413791b6119321b0e38a47b8db78ec67e042ee Mon Sep 17 00:00:00 2001 From: xFrednet Date: Fri, 6 Mar 2026 16:22:32 +0100 Subject: [PATCH 3/3] Ensure that the pre-freeze hook only runs once --- Include/object.h | 1 + Python/immutability.c | 51 ++++++++++++++++++++++++++++--------------- 2 files changed, 35 insertions(+), 17 deletions(-) diff --git a/Include/object.h b/Include/object.h index 5e5929e19eeb39..fab3a75d75a170 100644 --- a/Include/object.h +++ b/Include/object.h @@ -642,6 +642,7 @@ given type object has a specified feature. #if defined(Py_GIL_DISABLED) && defined(Py_DEBUG) #define _Py_TYPE_REVEALED_FLAG (1 << 3) #endif +#define _Py_PREFREEZE_RAN_FLAG (1 << 8) #define Py_CONSTANT_NONE 0 #define Py_CONSTANT_FALSE 1 diff --git a/Python/immutability.c b/Python/immutability.c index 0da016a1f2008d..35fc6e6932e6ea 100644 --- a/Python/immutability.c +++ b/Python/immutability.c @@ -1570,23 +1570,9 @@ void _Py_RefcntAdd_Immutable(PyObject *op, Py_ssize_t increment) // Macro that jumps to error, if the expression `x` does not succeed. #define SUCCEEDS(x) { do { int r = (x); if (r != 0) goto error; } while (0); } -static int run_pre_freeze_hook(PyObject* obj) { - PyObject *attr = NULL; - - // Skip Python-level hook lookup for type objects. For classes, - // `__pre_freeze__` resolves to an unbound function and calling it as a - // normal bound method would fail with a missing 'self' argument. - if (PyType_Check(obj)) { - return 0; - } - - // 0. Check if the pre-freeze hook already ran for this object - // TODO(immutable): Remember called hooks and return early - - // 1. Pre-freeze hooks are never called for shallow immutable objects - // TODO(immutable): Return early for shallow immutable. - +static int _run_pre_freeze_hook(struct _Py_immutability_state *imm_state, PyObject* obj) { // 1. Check for the `__pre_freeze__` name + PyObject *attr = NULL; int res = PyObject_GetOptionalAttr(obj, &_Py_ID(__pre_freeze__), &attr); if (res == -1) { return -1; @@ -1617,6 +1603,37 @@ static int run_pre_freeze_hook(PyObject* obj) { return 0; } +static int check_pre_freeze_hook(struct _Py_immutability_state *imm_state, PyObject* obj) { + // Skip Python-level hook lookup for type objects. For classes, + // `__pre_freeze__` resolves to an unbound function and calling it as a + // normal bound method would fail with a missing 'self' argument. + if (PyType_Check(obj)) { + return 0; + } + + // Pre-freeze hooks are never called for shallow immutable objects + if (is_shallow_immutable(imm_state, obj)) { + return 0; + } + + // Check if the pre-freeze hook already ran for this object +#if SIZEOF_VOID_P > 4 + if ((obj->ob_flags & _Py_PREFREEZE_RAN_FLAG) != 0) { + 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 + // the pre-freeze hook in case the pre-freeze hook reenters to prevent + // an infinite loop. + obj->ob_flags |= _Py_PREFREEZE_RAN_FLAG; + + // Run the pre-freeze hook if it's present. + return _run_pre_freeze_hook(imm_state, obj); +} + static int traverse_freeze(PyObject* obj, struct FreezeState* freeze_state) { // WARNING @@ -1796,7 +1813,7 @@ int _PyImmutability_Freeze(PyObject* obj) SUCCEEDS(check_freezable(imm_state, item, &freeze_state)); // Call the pre-freeze hook if one is present - SUCCEEDS(run_pre_freeze_hook(item)); + SUCCEEDS(check_pre_freeze_hook(imm_state, item)); // If the pre-freeze hook turned the object immutable, we want to skip it. if (_Py_IsImmutable(item)) {