Support for extension codes. (By accident I checked in the tests first.)
diff --git a/Lib/copy_reg.py b/Lib/copy_reg.py
index 8a3550a..1276564 100644
--- a/Lib/copy_reg.py
+++ b/Lib/copy_reg.py
@@ -70,3 +70,62 @@
return _reconstructor, args, dict
else:
return _reconstructor, args
+
+# A registry of extension codes. This is an ad-hoc compression
+# mechanism. Whenever a global reference to <module>, <name> is about
+# to be pickled, the (<module>, <name>) tuple is looked up here to see
+# if it is a registered extension code for it. Extension codes are
+# universal, so that the meaning of a pickle does not depend on
+# context. (There are also some codes reserved for local use that
+# don't have this restriction.) Codes are positive ints; 0 is
+# reserved.
+
+extension_registry = {} # key -> code
+inverted_registry = {} # code -> key
+extension_cache = {} # code -> object
+
+def add_extension(module, name, code):
+ """Register an extension code."""
+ code = int(code)
+ if not 1 <= code < 0x7fffffff:
+ raise ValueError, "code out of range"
+ key = (module, name)
+ if (extension_registry.get(key) == code and
+ inverted_registry.get(code) == key):
+ return # Redundant registrations are benign
+ if key in extension_registry:
+ raise ValueError("key %s is already registered with code %s" %
+ (key, extension_registry[key]))
+ if code in inverted_registry:
+ raise ValueError("code %s is already in use for key %s" %
+ (code, inverted_registry[code]))
+ extension_registry[key] = code
+ inverted_registry[code] = key
+
+def remove_extension(module, name, code):
+ """Unregister an extension code. For testing only."""
+ key = (module, name)
+ if (extension_registry.get(key) != code or
+ inverted_registry.get(code) != key):
+ raise ValueError("key %s is not registered with code %s" %
+ (key, code))
+ del extension_registry[key]
+ del inverted_registry[code]
+ if code in extension_cache:
+ del extension_cache[code]
+
+def clear_extension_cache():
+ extension_cache.clear()
+
+# Standard extension code assignments
+
+# Reserved ranges
+
+# First Last Count Purpose
+# 1 127 127 Reserved for Python standard library
+# 128 191 64 Reserved for Zope 3
+# 192 239 48 Reserved for 3rd parties
+# 240 255 16 Reserved for private use (will never be assigned)
+# 256 Inf Inf Reserved for future assignment
+
+# Extension codes are assigned by the Python Software Foundation.
diff --git a/Lib/pickle.py b/Lib/pickle.py
index 863702d..d98bcd5 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -28,6 +28,7 @@
from types import *
from copy_reg import dispatch_table, _reconstructor
+from copy_reg import extension_registry, inverted_registry, extension_cache
import marshal
import sys
import struct
@@ -395,16 +396,28 @@
self.memoize(obj)
if isinstance(obj, list):
- write(MARK)
- for x in obj:
- save(x)
- write(APPENDS)
+ n = len(obj)
+ if n > 1:
+ write(MARK)
+ for x in obj:
+ save(x)
+ write(APPENDS)
+ elif n == 1:
+ save(obj[0])
+ write(APPEND)
elif isinstance(obj, dict):
- write(MARK)
- for k, v in obj.iteritems():
+ n = len(obj)
+ if n > 1:
+ write(MARK)
+ for k, v in obj.iteritems():
+ save(k)
+ save(v)
+ write(SETITEMS)
+ elif n == 1:
+ k, v = obj.items()[0]
save(k)
save(v)
- write(SETITEMS)
+ write(SETITEM)
getstate = getattr(obj, "__getstate__", None)
if getstate:
@@ -420,6 +433,8 @@
getstate = None
if not getstate:
state = getattr(obj, "__dict__", None)
+ if not state:
+ state = None
# If there are slots, the state becomes a tuple of two
# items: the first item the regular __dict__ or None, and
# the second a dict mapping slot names to slot values
@@ -703,7 +718,7 @@
dispatch[InstanceType] = save_inst
- def save_global(self, obj, name = None):
+ def save_global(self, obj, name=None, pack=struct.pack):
write = self.write
memo = self.memo
@@ -729,6 +744,18 @@
"Can't pickle %r: it's not the same object as %s.%s" %
(obj, module, name))
+ if self.proto >= 2:
+ code = extension_registry.get((module, name))
+ if code:
+ assert code > 0
+ if code <= 0xff:
+ write(EXT1 + chr(code))
+ elif code <= 0xffff:
+ write(EXT2 + chr(code&0xff) + chr(code>>8))
+ else:
+ write(EXT4 + pack("<i", code))
+ return
+
write(GLOBAL + module + '\n' + name + '\n')
self.memoize(obj)
@@ -1081,6 +1108,34 @@
self.append(klass)
dispatch[GLOBAL] = load_global
+ def load_ext1(self):
+ code = ord(self.read(1))
+ self.get_extension(code)
+ dispatch[EXT1] = load_ext1
+
+ def load_ext2(self):
+ code = mloads('i' + self.read(2) + '\000\000')
+ self.get_extension(code)
+ dispatch[EXT2] = load_ext2
+
+ def load_ext4(self):
+ code = mloads('i' + self.read(4))
+ self.get_extension(code)
+ dispatch[EXT4] = load_ext4
+
+ def get_extension(self, code):
+ nil = []
+ obj = extension_cache.get(code, nil)
+ if obj is not nil:
+ self.append(obj)
+ return
+ key = inverted_registry.get(code)
+ if not key:
+ raise ValueError("unregistered extension code %d" % code)
+ obj = self.find_class(*key)
+ extension_cache[code] = obj
+ self.append(obj)
+
def find_class(self, module, name):
# Subclasses may override this
__import__(module)