Raise PicklingError when __reduce__() fails, and
add memoize() helper function to update the memo.
The first element of the tuple returned by __reduce__() must be a
callable. If it isn't the Unpickler will raise an error. Catch this
error in the pickler and raise the error there.
The memoize() helper also has a comment explaining how the memo
works. So methods can't use memoize() because the write funny codes.
diff --git a/Lib/pickle.py b/Lib/pickle.py
index a8b2b51..981cfe1 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -167,6 +167,22 @@
self.save(object)
self.write(STOP)
+ def memoize(self, obj):
+ """Store an object in the memo."""
+
+ # The memo is a dictionary mapping object ids to 2-tuples
+ # that contains the memo value and the object being memoized.
+ # The memo value is written to the pickle and will become
+ # the key in the Unpickler's memo. The object is stored in the
+ # memo so that transient objects are kept alive during pickling.
+
+ # The use of the memo length as the memo value is just a convention.
+ # The only requirement is that the memo values by unique.
+ d = id(obj)
+ memo_len = len(self.memo)
+ self.write(self.put(memo_len))
+ self.memo[d] = memo_len, obj
+
def put(self, i):
if self.bin:
s = mdumps(i)[1:]
@@ -280,11 +296,15 @@
self.save(pid)
self.write(BINPERSID)
- def save_reduce(self, callable, arg_tup, state = None):
+ def save_reduce(self, acallable, arg_tup, state = None):
write = self.write
save = self.save
- save(callable)
+ if not callable(acallable):
+ raise PicklingError("__reduce__() must return callable as "
+ "first argument, not %s" % `acallable`)
+
+ save(acallable)
save(arg_tup)
write(REDUCE)
@@ -340,9 +360,6 @@
dispatch[FloatType] = save_float
def save_string(self, object):
- d = id(object)
- memo = self.memo
-
if self.bin:
l = len(object)
s = mdumps(l)[1:]
@@ -352,16 +369,10 @@
self.write(BINSTRING + s + object)
else:
self.write(STRING + `object` + '\n')
-
- memo_len = len(memo)
- self.write(self.put(memo_len))
- memo[d] = (memo_len, object)
+ self.memoize(object)
dispatch[StringType] = save_string
def save_unicode(self, object):
- d = id(object)
- memo = self.memo
-
if self.bin:
encoding = object.encode('utf-8')
l = len(encoding)
@@ -371,17 +382,12 @@
object = object.replace("\\", "\\u005c")
object = object.replace("\n", "\\u000a")
self.write(UNICODE + object.encode('raw-unicode-escape') + '\n')
-
- memo_len = len(memo)
- self.write(self.put(memo_len))
- memo[d] = (memo_len, object)
+ self.memoize(object)
dispatch[UnicodeType] = save_unicode
if StringType == UnicodeType:
# This is true for Jython
def save_string(self, object):
- d = id(object)
- memo = self.memo
unicode = object.isunicode()
if self.bin:
@@ -404,14 +410,10 @@
self.write(UNICODE + object + '\n')
else:
self.write(STRING + `object` + '\n')
-
- memo_len = len(memo)
- self.write(self.put(memo_len))
- memo[d] = (memo_len, object)
+ self.memoize(object)
dispatch[StringType] = save_string
def save_tuple(self, object):
-
write = self.write
save = self.save
memo = self.memo
@@ -434,6 +436,7 @@
memo_len = len(memo)
self.write(TUPLE + self.put(memo_len))
memo[d] = (memo_len, object)
+
dispatch[TupleType] = save_tuple
def save_empty_tuple(self, object):
@@ -451,9 +454,7 @@
else:
write(MARK + LIST)
- memo_len = len(memo)
- write(self.put(memo_len))
- memo[d] = (memo_len, object)
+ self.memoize(object)
using_appends = (self.bin and (len(object) > 1))
@@ -471,20 +472,15 @@
dispatch[ListType] = save_list
def save_dict(self, object):
- d = id(object)
-
write = self.write
save = self.save
- memo = self.memo
if self.bin:
write(EMPTY_DICT)
else:
write(MARK + DICT)
- memo_len = len(memo)
- self.write(self.put(memo_len))
- memo[d] = (memo_len, object)
+ self.memoize(object)
using_setitems = (self.bin and (len(object) > 1))
@@ -529,6 +525,8 @@
for arg in args:
save(arg)
+ # This method does not use memoize() so that it can handle
+ # the special case for non-binary mode.
memo_len = len(memo)
if self.bin:
write(OBJ + self.put(memo_len))