Cleaning up the immutability code

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

Author: matajoh

Include/internal/pycore_object.h

@@ -101,7 +101,7 @@ static inline void _Py_SetImmutable(PyObject *op)
 #define _Py_SetImmutable(op) _Py_SetImmutable(_PyObject_CAST(op))
 
 // Check whether an object is writeable.
-// Note that during runtime finalization, all objects must be mutable
+// This check will always succeed during runtime finalization.
 #define Py_CHECKWRITE(op) ((op) && (!_Py_IsImmutable(op) || _Py_IsFinalizing()))
 #define Py_REQUIREWRITE(op, msg) {if (Py_CHECKWRITE(op)) { _PyObject_ASSERT_FAILED_MSG(op, msg); }}
 

Lib/test/test_freeze.py

@@ -1,4 +1,6 @@
 import unittest
+from collections import deque
+from array import array
 
 # This is a canary to check that global variables are not made immutable
 # when others are made immutable
@@ -97,6 +99,83 @@ def test_reverse(self):
         with self.assertRaises(NotWriteableError):
             self.obj.reverse()
 
+    def test_inplace_repeat(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj *= 2
+
+    def test_inplace_concat(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj += [TestList.C()]
+
+    def test_clear(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.clear()
+
+    def test_sort(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.sort()
+
+
+class TestDeque(BaseObjectTest):
+    class C:
+        pass
+
+    def __init__(self, *args, **kwargs):
+        obj = deque([self.C(), self.C(), 1, "two", None])
+        BaseObjectTest.__init__(self, *args, obj=obj, **kwargs)
+
+    def test_set_item(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj[0] = None
+
+    def test_set_slice(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj[1:3] = [None, None]
+
+    def test_append(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.append(TestList.C())
+
+    def test_appendleft(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.appendleft(TestList.C())
+
+    def test_extend(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.extend([TestList.C()])
+
+    def test_extendleft(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.extendleft([TestList.C()])
+
+    def test_insert(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.insert(0, TestList.C())
+
+    def test_pop(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.pop()
+
+    def test_popleft(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.popleft()
+
+    def test_remove(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.remove(1)
+
+    def test_inplace_repeat(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj *= 2
+
+    def test_inplace_concat(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj += [TestList.C()]
+
+    def test_reverse(self):
+        with self.assertRaises(NotWriteableError):
+            self.obj.reverse()
+
     def test_clear(self):
         with self.assertRaises(NotWriteableError):
             self.obj.clear()

Objects/listobject.c

@@ -774,6 +774,12 @@ list_inplace_repeat(PyListObject *self, Py_ssize_t n)
     if (input_size > PY_SSIZE_T_MAX / n) {
         return PyErr_NoMemory();
     }
+
+    if(!Py_CHECKWRITE(self))
+    {
+        return PyErr_WriteToImmutable(self);
+    }
+
     Py_ssize_t output_size = input_size * n;
 
     if (list_resize(self, output_size) < 0)

Python/immutability.c

@@ -8,19 +8,32 @@
 #include "pycore_immutability.h"
 
 
+static int push(PyObject* s, PyObject* item){
+    if(item == NULL){
+        return 0;
+    }
+
+    if(!PyList_Check(s)){
+        PyErr_SetString(PyExc_TypeError, "Expected a list");
+        return -1;
+    }
+
+    return _PyList_AppendTakeRef(_PyList_CAST(s), item);
+}
+
 static PyObject* pop(PyObject* s){
+    PyObject* item;
     Py_ssize_t size = PyList_Size(s);
     if(size == 0){
         return NULL;
     }
 
-    PyObject* item = PyList_GetItem(s, size - 1);
+    item = PyList_GetItem(s, size - 1);
     if(item == NULL){
         return NULL;
     }
 
     Py_INCREF(item);
-
     if(PyList_SetSlice(s, size - 1, size, NULL)){
         Py_DECREF(item);
         return NULL;
@@ -35,7 +48,7 @@ static bool is_c_wrapper(PyObject* obj){
 
 #define _Py_VISIT_FUNC_ATTR(attr, frontier) do { \
     if(attr != NULL && !_Py_IsImmutable(attr)){ \
-        if(PyList_Append(frontier, attr)){ \
+        if(push((frontier), (attr))){ \
             return PyErr_NoMemory(); \
         } \
     } \
@@ -118,7 +131,7 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
     }
 
     f_ptr = f->func_code;
-    if(PyList_Append(f_stack, f_ptr)){
+    if(push(f_stack, f_ptr)){
         goto nomemory;
     }
 
@@ -174,7 +187,7 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
                 _PyDict_SetKeyImmutable((PyDictObject*)module_dict, name);
 
                 if(!_Py_IsImmutable(value)){
-                    if(PyList_Append(frontier, value)){
+                    if(push(frontier, value)){
                         goto nomemory;
                     }
                 }
@@ -188,11 +201,11 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
                 if(PyCode_Check(value)){
                     _Py_SetImmutable(value);
 
-                    if(PyList_Append(f_stack, value)){
+                    if(push(f_stack, value)){
                         goto nomemory;
                     }
                 }else{
-                    if(PyList_Append(frontier, value)){
+                    if(push(frontier, value)){
                         goto nomemory;
                     }
                 }
@@ -250,14 +263,14 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
         }
     }
 
-    if(PyList_Append(frontier, frozen_globals)){
+    if(push(frontier, frozen_globals)){
         goto nomemory;
     }
 
     f->func_globals = frozen_globals;
     Py_DECREF(globals);
 
-    if(PyList_Append(frontier, frozen_builtins)){
+    if(push(frontier, frozen_builtins)){
         goto nomemory;
     }
 
@@ -276,7 +289,7 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
 static int freeze_visit(PyObject* obj, void* frontier)
 {
     if(!_Py_IsImmutable(obj)){
-        if(PyList_Append((PyObject*)frontier, obj)){
+        if(push(frontier, obj)){
             PyErr_NoMemory();
             return -1;
         }
@@ -287,44 +300,49 @@ static int freeze_visit(PyObject* obj, void* frontier)
 
 PyObject* _Py_Freeze(PyObject* obj)
 {
+    PyObject* frontier = NULL;
+    PyObject* frozen_importlib = NULL;
+    PyObject* blocking_on = NULL;
+    PyObject* module_locks = NULL;
+    PyObject* result = Py_None;
+
     if(_Py_IsImmutable(obj)){
-        Py_RETURN_NONE;
+        return result;
     }
 
-    PyObject* frontier = PyList_New(0);
+    frontier = PyList_New(0);
     if(frontier == NULL){
-        return PyErr_NoMemory();
+        result = PyErr_NoMemory();
+        goto cleanup;
     }
 
-    if(PyList_Append(frontier, obj)){
-        Py_DECREF(frontier);
-        return PyErr_NoMemory();
+    if(push(frontier, obj)){
+        result = PyErr_NoMemory();
+        goto cleanup;
     }
 
-    PyObject* frozen_importlib = PyImport_ImportModule("_frozen_importlib");
+    frozen_importlib = PyImport_ImportModule("_frozen_importlib");
     if(frozen_importlib == NULL){
-        Py_DECREF(frontier);
-        return NULL;
+        result = NULL;
+        goto cleanup;
     }
 
-    PyObject* blocking_on = PyObject_GetAttrString(frozen_importlib, "_blocking_on");
+    blocking_on = PyObject_GetAttrString(frozen_importlib, "_blocking_on");
     if(blocking_on == NULL){
-        Py_DECREF(frozen_importlib);
-        Py_DECREF(frontier);
-        return NULL;
+        result = NULL;
+        goto cleanup;
     }
 
-    PyObject* module_locks = PyObject_GetAttrString(frozen_importlib, "_module_locks");
+    module_locks = PyObject_GetAttrString(frozen_importlib, "_module_locks");
     if(module_locks == NULL){
-        Py_DECREF(blocking_on);
-        Py_DECREF(frozen_importlib);
-        Py_DECREF(frontier);
-        return NULL;
+        result = NULL;
+        goto cleanup;
     }
 
-    Py_DECREF(frozen_importlib);
-
     while(PyList_Size(frontier) != 0){
+        PyTypeObject* type;
+        PyObject* type_op;
+        traverseproc traverse;
         PyObject* item = pop(frontier);
 
         if(item == blocking_on ||
@@ -334,9 +352,8 @@ PyObject* _Py_Freeze(PyObject* obj)
             continue;
         }
 
-        PyTypeObject* type = Py_TYPE(item);
-        traverseproc traverse;
-        PyObject* type_op = NULL;
+        type = Py_TYPE(item);
+        type_op = NULL;
 
         if(_Py_IsImmutable(item)){
             continue;
@@ -350,42 +367,37 @@ 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);
-                Py_DECREF(frontier);
-                return err;
+            result = walk_function(item, frontier);
+            if(!Py_IsNone(result)){
+                goto cleanup;
             }
         }
         else
         {
             traverse = type->tp_traverse;
             if(traverse != NULL){
                 if(traverse(item, (visitproc)freeze_visit, frontier)){
-                    Py_DECREF(blocking_on);
-                    Py_DECREF(module_locks);
-                    Py_DECREF(frontier);
-                    return NULL;
+                    result = NULL;
+                    goto cleanup;
                 }
             }
         }
 
         type_op = _PyObject_CAST(item->ob_type);
         if (!_Py_IsImmutable(type_op)){
-            if (PyList_Append(frontier, type_op))
+            if(push(frontier, type_op))
             {
-                Py_DECREF(blocking_on);
-                Py_DECREF(module_locks);
-                Py_DECREF(frontier);
-                return PyErr_NoMemory();
+                result = PyErr_NoMemory();
+                goto cleanup;
             }
         }
     }
 
-    Py_DECREF(blocking_on);
-    Py_DECREF(module_locks);
-    Py_DECREF(frontier);
+cleanup:
+    Py_XDECREF(blocking_on);
+    Py_XDECREF(module_locks);
+    Py_XDECREF(frozen_importlib);
+    Py_XDECREF(frontier);
 
-    Py_RETURN_NONE;
+    return result;
 }