diff options
| -rw-r--r-- | Lib/copy_reg.py | 59 | ||||
| -rw-r--r-- | Lib/pickle.py | 71 |
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) |