Support for extension codes. (By accident I checked in the tests first.)
This commit is contained in:
parent
0322d0ff13
commit
255f3ee0a5
|
@ -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.
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue