Support for extension codes. (By accident I checked in the tests first.)

2 files changed, 122 insertions, 8 deletions

diff --git a/Lib/copy_reg.py b/Lib/copy_reg.py
index 8a3550aad3..1276564b32 100644
--- a/Lib/copy_reg.py
+++ b/Lib/copy_reg.py
@@ -70,3 +70,62 @@ def _reduce(self):
         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 863702dc9f..d98bcd5076 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -28,6 +28,7 @@ __version__ = "$Revision$"       # Code version
 
 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 @@ class Pickler:
         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 @@ class Pickler:
                 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 @@ class Pickler:
 
     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 @@ class Pickler:
                     "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 @@ class Unpickler:
         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)