pickle supports frozendict

Co-Authored-by: Marco Sulla <[email protected]>

Author: vstinner

Lib/pickle.py

@@ -185,6 +185,7 @@ def __init__(self, value):
 BYTEARRAY8       = b'\x96'  # push bytearray
 NEXT_BUFFER      = b'\x97'  # push next out-of-band buffer
 READONLY_BUFFER  = b'\x98'  # make top of stack readonly
+FROZENDICT       = b'\x99'
 
 __all__.extend(x for x in dir() if x.isupper() and not x.startswith('_'))
 
@@ -1064,6 +1065,31 @@ def save_dict(self, obj):
 
     dispatch[dict] = save_dict
 
+    def save_frozendict(self, obj):
+        save = self.save
+        write = self.write
+
+        if self.proto < 5:
+            self.save_reduce(frozendict, (dict(obj),), obj=obj)
+            return
+
+        write(MARK)
+        for k, v in obj.items():
+            save(k)
+            save(v)
+
+        if id(obj) in self.memo:
+            # If the object is already in the memo, this means it is
+            # recursive. In this case, throw away everything we put on the
+            # stack, and fetch the object back from the memo.
+            write(POP_MARK + self.get(self.memo[id(obj)][0]))
+            return
+
+        write(FROZENDICT)
+        self.memoize(obj)
+
+    dispatch[frozendict] = save_frozendict
+
     def _batch_setitems(self, items, obj):
         # Helper to batch up SETITEMS sequences; proto >= 1 only
         save = self.save
@@ -1589,6 +1615,13 @@ def load_dict(self):
         self.append(d)
     dispatch[DICT[0]] = load_dict
 
+    def load_frozendict(self):
+        items = self.pop_mark()
+        d = frozendict({items[i]: items[i+1] for i in range(0, len(items), 2)})
+        self.append(d)
+
+    dispatch[FROZENDICT[0]] = load_frozendict
+
     # INST and OBJ differ only in how they get a class object.  It's not
     # only sensible to do the rest in a common routine, the two routines
     # previously diverged and grew different bugs.

Lib/pickletools.py

@@ -1035,6 +1035,11 @@ def __repr__(self):
     obtype=dict,
     doc="A Python dict object.")
 
+pyfrozendict = StackObject(
+    name="frozendict",
+    obtype=frozendict,
+    doc="A Python frozendict object.")
+
 pyset = StackObject(
     name="set",
     obtype=set,
@@ -1384,6 +1389,23 @@ def __init__(self, name, code, arg,
       proto=5,
       doc="Make an out-of-band buffer object read-only."),
 
+    I(name='FROZENDICT',
+      code='\x99',
+      arg=None,
+      stack_before=[markobject, stackslice],
+      stack_after=[pyfrozendict],
+      proto=5,
+      doc="""Build a frozendict out of the topmost stack slice, after markobject.
+
+      All the stack entries following the topmost markobject are placed into
+      a single Python dict, which single dict object replaces all of the
+      stack from the topmost markobject onward.  The stack slice alternates
+      key, value, key, value, ....  For example,
+
+      Stack before: ... markobject 1 2 3 'abc'
+      Stack after:  ... {1: 2, 3: 'abc'}
+      """),
+
     # Ways to spell None.
 
     I(name='NONE',

Modules/_pickle.c

@@ -136,7 +136,8 @@ enum opcode {
     /* Protocol 5 */
     BYTEARRAY8       = '\x96',
     NEXT_BUFFER      = '\x97',
-    READONLY_BUFFER  = '\x98'
+    READONLY_BUFFER  = '\x98',
+    FROZENDICT       = '\x99',
 };
 
 enum {
@@ -592,6 +593,36 @@ Pdata_poplist(Pdata *self, Py_ssize_t start)
     return list;
 }
 
+static PyObject *
+Pdata_poplist2(PickleState *state, Pdata *self, Py_ssize_t start)
+{
+    if (start < self->fence) {
+        Pdata_stack_underflow(state, self);
+        return NULL;
+    }
+
+    Py_ssize_t len = (Py_SIZE(self) - start) >> 1;
+
+    PyObject *list = PyList_New(len);
+    if (list == NULL) {
+        return NULL;
+    }
+
+    for (Py_ssize_t i = start, j = 0; j < len; i+=2, j++) {
+        PyObject *subtuple = PyTuple_New(2);
+        if (subtuple == NULL) {
+            return NULL;
+        }
+
+        PyTuple_SET_ITEM(subtuple, 0, self->data[i]);
+        PyTuple_SET_ITEM(subtuple, 1, self->data[i+1]);
+        PyList_SET_ITEM(list, j, subtuple);
+    }
+
+    Py_SET_SIZE(self, start);
+    return list;
+}
+
 typedef struct {
     PyObject *me_key;
     Py_ssize_t me_value;
@@ -3445,6 +3476,64 @@ save_dict(PickleState *state, PicklerObject *self, PyObject *obj)
     return status;
 }
 
+static int
+save_frozendict(PickleState *state, PicklerObject *self, PyObject *obj)
+{
+    const char mark_op = MARK;
+    const char frozendict_op = FROZENDICT;
+
+    if (self->fast && !fast_save_enter(self, obj)) {
+        return -1;
+    }
+
+    if (self->proto < 4) {
+        PyObject *items = PyDict_Items(obj);
+        if (items == NULL) {
+            return -1;
+        }
+
+        PyObject *reduce_value;
+        reduce_value = Py_BuildValue("(O(O))", (PyObject*)&PyFrozenDict_Type,
+                                     items);
+        Py_DECREF(items);
+        if (reduce_value == NULL) {
+            return -1;
+        }
+
+        /* save_reduce() will memoize the object automatically. */
+        int status = save_reduce(state, self, reduce_value, obj);
+        Py_DECREF(reduce_value);
+        return status;
+    }
+
+    if (_Pickler_Write(self, &mark_op, 1) < 0) {
+        return -1;
+    }
+
+    PyObject *key = NULL, *value = NULL;
+    Py_ssize_t pos = 0;
+    while (PyDict_Next(obj, &pos, &key, &value)) {
+        int res = save(state, self, key, 0);
+        if (res < 0) {
+            return -1;
+        }
+
+        res = save(state, self, value, 0);
+        if (res < 0) {
+            return -1;
+        }
+    }
+
+    if (_Pickler_Write(self, &frozendict_op, 1) < 0) {
+        return -1;
+    }
+
+    if (memo_put(state, self, obj) < 0) {
+        return -1;
+    }
+    return 0;
+}
+
 static int
 save_set(PickleState *state, PicklerObject *self, PyObject *obj)
 {
@@ -4411,6 +4500,10 @@ save(PickleState *st, PicklerObject *self, PyObject *obj, int pers_save)
         status = save_dict(st, self, obj);
         goto done;
     }
+    else if (type == &PyFrozenDict_Type) {
+        status = save_frozendict(st, self, obj);
+        goto done;
+    }
     else if (type == &PySet_Type) {
         status = save_set(st, self, obj);
         goto done;
@@ -5831,6 +5924,30 @@ load_dict(PickleState *st, UnpicklerObject *self)
     return 0;
 }
 
+
+static int
+load_frozendict(PickleState *st, UnpicklerObject *self)
+{
+    Py_ssize_t i = marker(st, self);
+    if (i < 0) {
+        return -1;
+    }
+
+    PyObject *items = Pdata_poplist2(st, self->stack, i);
+    if (items == NULL) {
+        return -1;
+    }
+
+    PyObject *frozendict = PyFrozenDict_New(items);
+    Py_DECREF(items);
+    if (frozendict == NULL) {
+        return -1;
+    }
+
+    PDATA_PUSH(self->stack, frozendict, -1);
+    return 0;
+}
+
 static int
 load_frozenset(PickleState *state, UnpicklerObject *self)
 {
@@ -6946,6 +7063,7 @@ load(PickleState *st, UnpicklerObject *self)
         OP(LIST, load_list)
         OP(EMPTY_DICT, load_empty_dict)
         OP(DICT, load_dict)
+        OP(FROZENDICT, load_frozendict)
         OP(EMPTY_SET, load_empty_set)
         OP(ADDITEMS, load_additems)
         OP(FROZENSET, load_frozenset)