*** empty log message ***

lib/Components/cPickle/pickle.py

-#     $Id: pickle.py,v 1.4 1997/01/03 21:21:10 chris Exp $
+#     $Id: pickle.py,v 1.5 1997/02/10 22:08:32 cici Exp $
 #
 #     Copyright 
 #
@@ -184,6 +184,7 @@ I have no answers.  Garbage Collection may also become a problem here.)
 __version__ = "1.6"                     # Code version
 
 from types import *
+from copy_reg import *
 import string, marshal
 
 format_version = "1.2"                  # File format version we write
@@ -193,57 +194,46 @@ mdumps = marshal.dumps
 mloads = marshal.loads
 
 PicklingError = "pickle.PicklingError"
-
-AtomicTypes = [NoneType, IntType, FloatType, StringType]
-
-def safe(object):
-    t = type(object)
-    if t in AtomicTypes:
-	return 1
-    if t is TupleType:
-	for item in object:
-	    if not safe(item): return 0
-	return 1
-    return 0
+UnpicklingError = "pickle.UnpicklingError"
 
 MARK            = '('
+STOP            = '.'
 POP             = '0'
 DUP             = '2'
-STOP = '.'
-TUPLE = 't'
-LIST = 'l'
-DICT = 'd'
-INST = 'i'
-OBJ = 'o'
-CLASS = 'c'
-GET = 'g'
-BINGET = 'h'
-PUT = 'p'
-BINPUT = 'q'
-APPEND = 'a'
-SETITEM = 's'
-BUILD = 'b'
-NONE = 'N'
+FLOAT           = 'F'
 INT             = 'I'
 BININT          = 'J'
 BININT1         = 'K'
-BININT2 = 'M'
-BININT3 = 'O'
 LONG            = 'L'
-FLOAT = 'F'
+BININT2         = 'M'
+NONE            = 'N'
+BININT8         = 'O'
+PERSID          = 'P'
+BINPERSID       = 'Q'
+REDUCE          = 'R'
 STRING          = 'S'
 BINSTRING       = 'T'
 SHORT_BINSTRING = 'U'
-PERSID = 'P'
-
-AtomicKeys = [NONE, INT, LONG, FLOAT, STRING]
-AtomicMap = {
-        NoneType: NONE,
-        IntType: INT,
-        LongType: LONG,
-        FloatType: FLOAT,
-        StringType: STRING,
-}
+APPEND          = 'a'
+BUILD           = 'b'
+GLOBAL          = 'c'
+DICT            = 'd'
+EMPTY_DICT      = '}'
+APPENDS         = 'e'
+GET             = 'g'
+BINGET          = 'h'
+INST            = 'i'
+LONG_BINGET     = 'j'
+LIST            = 'l'
+EMPTY_LIST      = ']'
+OBJ             = 'o'
+PUT             = 'p'
+BINPUT          = 'q'
+LONG_BINPUT     = 'r'
+SETITEM         = 's'
+TUPLE           = 't'
+EMPTY_TUPLE     = ')'
+SETITEMS        = 'u'
 
 class Pickler:
 
@@ -256,67 +246,149 @@ class Pickler:
         self.save(object)
         self.write(STOP)
 
-    def put_str(self, i):
-	if (self.bin and (i < 256)):
-            return BINPUT + chr(i)
+    def dump_special(self, callable, args, state = None):
+	if (type(args) is not TupleType):
+            raise PicklingError, "Second argument to dump_special " \
+                                 "must be a tuple"
+
+        self.save_reduce(callable, args, state)
+        self.write(STOP)
+
+    def put(self, i):
+        if (self.bin):
+            s = mdumps(i)[1:]
+            if (i < 256):
+                return BINPUT + s[0]
+
+            return LONG_BINPUT + s
 
         return PUT + `i` + '\n'
 
-    def get_str(self, i):
-	if (self.bin and (i < 256)):
-            return BINGET + chr(i)
+    def get(self, i):
+        if (self.bin):
+            s = mdumps(i)[1:]
+
+            if (i < 256):
+                return BINGET + s[0]
+
+            return LONG_BINGET + s
 
         return GET + `i` + '\n'
         
-    def save(self, object):
+    def save(self, object, pers_save = 0):
         memo = self.memo
+
+        if (not pers_save):
 	    pid = self.persistent_id(object)
-	if pid:
-	    self.write(PERSID + str(pid) + '\n')
+	    if (pid is not None):
+                self.save_pers(pid)
                 return
+
         d = id(object)
+ 
+        t = type(object)
+        try:
+            f = self.quick_dispatch[t]
+        except KeyError:
             if memo.has_key(d):
-            self.write(self.get_str(memo[d]))
+                self.write(self.get(memo[d]))
                 return
-	t = type(object)
+
             try:
                 f = self.dispatch[t]
             except KeyError:
-	    if hasattr(object, '__class__'):
-		f = self.dispatch[InstanceType]
-	    else:
+                pid = self.inst_persistent_id(object)
+                if pid is not None:
+                    self.save_pers(pid)
+                    return
+
+                try:
+                    reduce = dispatch_table[t]
+                except KeyError:
+                    try:
+                        reduce = object.__reduce__
+                    except AttributeError:
                         raise PicklingError, \
                             "can't pickle %s objects" % `t.__name__`
+		    else:
+                        tup = reduce()
+                else:
+                    tup = reduce(object)
+
+                if (type(tup) is not TupleType):
+                    raise PicklingError, "Value returned by %s must be a " \
+                                         "tuple" % reduce
+
+                l = len(tup)
+   
+		if ((l != 2) and (l != 3)):
+                    raise PicklingError, "tuple returned by %s must contain " \
+                                         "only two or three elements" % reduce
+
+                callable = tup[0]
+                arg_tup  = tup[1]
+          
+                if (l > 2):
+                    state = tup[2]
+                else:
+                    state = None
+
+                if (type(arg_tup) is not TupleType):
+                    raise PicklingError, "Second element of tuple returned " \
+                                         "by %s must be a tuple" % reduce
+
+                self.save_reduce(callable, arg_tup, state) 
+                return
+
         f(self, object)
 
     def persistent_id(self, object):
         return None
 
+    def inst_persistent_id(self, object):
+        return None
+
+    def save_pers(self, pid):
+        if (not self.bin):
+            self.write(PERSID + str(pid) + '\n')
+        else:
+            self.save(pid, 1)
+            self.write(BINPERSID)
+
+    def save_reduce(self, callable, arg_tup, state = None):
+        write = self.write
+        save = self.save
+
+        save(callable)
+        save(arg_tup)
+        write(REDUCE)
+        
+	if (state is not None):
+            save(state)
+            write(BUILD)
+
     dispatch = {}
+    quick_dispatch = {}
 
     def save_none(self, object):
         self.write(NONE)
-    dispatch[NoneType] = save_none
+    quick_dispatch[NoneType] = save_none
 
     def save_int(self, object):
         if (self.bin):
             i = mdumps(object)[1:]
-	    if (i[-1] == '\000'):
-		if (i[-2] == '\000'):
+            if (i[-2:] == '\000\000'):
                 if (i[-3] == '\000'):
-			self.write(BININT3 + i[:-3])
+                    self.write(BININT1 + i[:-3])
                     return
 
                 self.write(BININT2 + i[:-2])
                 return
 
-  	        self.write(BININT1 + i[:-1])
-	        return
-
             self.write(BININT + i)
         else:
             self.write(INT + `object` + '\n')
-    dispatch[IntType] = save_int
+    quick_dispatch[IntType] = save_int
 
     def save_long(self, object):
         self.write(LONG + `object` + '\n')
@@ -324,7 +396,7 @@ class Pickler:
 
     def save_float(self, object):
         self.write(FLOAT + `object` + '\n')
-    dispatch[FloatType] = save_float
+    quick_dispatch[FloatType] = save_float
 
     def save_string(self, object):
         d = id(object)
@@ -332,15 +404,16 @@ class Pickler:
 
         if (self.bin):
             l = len(object)
+            s = mdumps(l)[1:]
             if (l < 256):
-	        self.write(SHORT_BINSTRING + chr(l) + object)
+                self.write(SHORT_BINSTRING + s[0] + object)
             else:
-	        self.write(BINSTRING + `l` + '\n' + object)
+                self.write(BINSTRING + s + object)
         else:
             self.write(STRING + `object` + '\n')
 
         memo_len = len(memo)
-        self.write(self.put_str(memo_len))
+        self.write(self.put(memo_len))
         memo[d] = memo_len
     dispatch[StringType] = save_string
 
@@ -350,27 +423,25 @@ class Pickler:
         write = self.write
         save  = self.save
         memo  = self.memo
-	has_key = self.memo.has_key
 
         write(MARK)
 
-	n = len(object)
-	for k in range(n):
-	    save(object[k])
+        for element in object:
+            save(element)
+
+        if (memo.has_key(d)):
+            write(POP * len(object))
+            write(self.get(memo[d]))
+            return
 
-	    if has_key(d):
-		# Saving object[k] has saved us!
-		while k >= 0:
-		    write(POP)
-		    k = k-1
-                write(self.get_str(memo[d]))
-		break
-	else:
         memo_len = len(memo)
-            self.write(TUPLE + self.put_str(memo_len))
+        self.write(TUPLE + self.put(memo_len))
         memo[d] = memo_len
     dispatch[TupleType] = save_tuple
 
+    def save_empty_tuple(self, object):
+        self.write(EMPTY_TUPLE)
+
     def save_list(self, object):
         d = id(object)
 
@@ -378,25 +449,25 @@ class Pickler:
         save  = self.save
         memo  = self.memo
 
-	write(MARK)
+        using_appends = (self.bin and (len(object) > 1))
 
-	n = len(object)
-	for k in range(n):
-	    item = object[k]
-	    if not safe(item):
-		break
-	    save(item)
+        if (using_appends):
+            write(MARK + LIST + MARK)
         else:
-	    k = n
+            write(MARK + LIST)
 
-	memo_len = len(memo)
-        write(LIST + self.put_str(memo_len))
-	memo[d] = memo_len
+        for element in object:
+            save(element)
   
-	for k in range(k, n):
-	    item = object[k]
-	    save(item)
+            if (not using_appends):
                 write(APPEND)
+
+        if (using_appends):
+            write(APPENDS)
+
+        memo_len = len(memo)
+        write(self.put(memo_len))
+        memo[d] = memo_len
     dispatch[ListType] = save_list
 
     def save_dict(self, object):
@@ -406,28 +477,27 @@ class Pickler:
         save  = self.save
         memo  = self.memo
 
-	write(MARK)
+        using_setitems = (self.bin and (len(object) > 1))
+
+        if (using_setitems):
+            write(MARK + DICT + MARK)
+        else:
+            write(MARK + DICT)
 
         items = object.items()
-	n = len(items)
-	for k in range(n):
-	    key, value = items[k]
-	    if not safe(key) or not safe(value):
-		break
+        for key, value in items:
             save(key)
             save(value)
-	else:
-	    k = n
+
+            if (not using_setitems):
+                write(SETITEM)
+
+        if (using_setitems):
+            write(SETITEMS)
 
         memo_len = len(memo)
-	self.write(DICT + self.put_str(memo_len))
+        self.write(self.put(memo_len))
         memo[d] = memo_len
-
-	for k in range(k, n):
-	    key, value = items[k]
-	    save(key)
-	    save(value)
-	    write(SETITEM)
     dispatch[DictionaryType] = save_dict
 
     def save_inst(self, object):
@@ -454,12 +524,12 @@ class Pickler:
 
         memo_len = len(memo)
         if (self.bin):
-            write(OBJ + self.put_str(memo_len))
+            write(OBJ + self.put(memo_len))
         else:
-            module = whichmodule(cls)
+            module = whichmodule(cls, cls.__name__)
             name = cls.__name__
             write(INST + module + '\n' + name + '\n' +
-                self.put_str(memo_len))
+                self.put(memo_len))
 
         memo[d] = memo_len
 
@@ -473,23 +543,27 @@ class Pickler:
         write(BUILD)
     dispatch[InstanceType] = save_inst
 
-    def save_class(self, object):
+    def save_global(self, object, name = None):
+        write = self.write
         memo = self.memo
-	d = id(object)
-	module = whichmodule(object)
+
+        if (name is None):
             name = object.__name__
 
-	memo_len = len(memo)
-	self.write(CLASS + module + '\n' + name + '\n' + 
-	      self.put_str(memo_len))
-	memo[d] = memo_len
+        module = whichmodule(object, name)
 
-    dispatch[ClassType] = save_class
+        memo_len = len(memo)
+        write(GLOBAL + module + '\n' + name + '\n' +
+            self.put(memo_len))
+        memo[id(object)] = memo_len
+    dispatch[ClassType] = save_global
+    dispatch[FunctionType] = save_global
+    dispatch[BuiltinFunctionType] = save_global
 
 
 classmap = {}
 
-def whichmodule(cls):
+def whichmodule(cls, clsname):
     """Figure out the module in which a class occurs.
 
     Search sys.modules for the module.
@@ -500,7 +574,7 @@ def whichmodule(cls):
     if classmap.has_key(cls):
         return classmap[cls]
     import sys
-    clsname = cls.__name__
+
     for name, module in sys.modules.items():
         if name != '__main__' and \
            hasattr(module, clsname) and \
@@ -550,6 +624,15 @@ class Unpickler:
         self.append(self.persistent_load(pid))
     dispatch[PERSID] = load_persid
 
+    def load_binpersid(self):
+        stack = self.stack
+         
+        pid = stack[-1]
+        del stack[-1]
+
+        self.append(self.persistent_load(pid))
+    dispatch[BINPERSID] = load_binpersid
+
     def load_none(self):
         self.append(None)
     dispatch[NONE] = load_none
@@ -563,17 +646,13 @@ class Unpickler:
     dispatch[BININT] = load_binint
 
     def load_binint1(self):
-        self.append(mloads('i' + self.read(3) + '\000'))
+        self.append(mloads('i' + self.read(1) + '\000\000\000'))
     dispatch[BININT1] = load_binint1
 
     def load_binint2(self):
         self.append(mloads('i' + self.read(2) + '\000\000'))
     dispatch[BININT2] = load_binint2
  
-    def load_binint3(self):
-        self.append(mloads('i' + self.read(1) + '\000\000\000'))
-    dispatch[BININT3] = load_binint3
-
     def load_long(self):
         self.append(string.atol(self.readline()[:-1], 0))
     dispatch[LONG] = load_long
@@ -588,12 +667,13 @@ class Unpickler:
     dispatch[STRING] = load_string
 
     def load_binstring(self):
-	self.append(self.read(string.atoi(self.readline()[:-1])))
+        len = mloads('i' + self.read(4))
+        self.append(self.read(len))
     dispatch[BINSTRING] = load_binstring
 
     def load_short_binstring(self):
-        read = self.read
-        self.append(read(ord(read(1))))
+        len = mloads('i' + self.read(1))
+        self.append(self.read(len))
     dispatch[SHORT_BINSTRING] = load_short_binstring
 
     def load_tuple(self):
@@ -624,6 +704,10 @@ class Unpickler:
         module = self.readline()[:-1]
         name = self.readline()[:-1]
         klass = self.find_class(module, name)
+	if (type(klass) is not ClassType):
+            raise SystemError, "Imported object %s from module %s is " \
+                               "not a class" % (name, module)
+
         value = apply(klass, args)
         self.append(value)
     dispatch[INST] = load_inst
@@ -639,13 +723,12 @@ class Unpickler:
         self.append(value)
     dispatch[OBJ] = load_obj                
 
-    def load_class(self):
+    def load_global(self):
         module = self.readline()[:-1]
         name = self.readline()[:-1]
         klass = self.find_class(module, name)
         self.append(klass)
-	return klass
-    dispatch[CLASS] = load_class
+    dispatch[GLOBAL] = load_global
 
     def find_class(self, module, name):
         env = {}
@@ -656,12 +739,30 @@ class Unpickler:
                   "Failed to import class %s from module %s" % \
                   (name, module)
         klass = env[name]
-	if type(klass) is BuiltinFunctionType:
-	    raise SystemError, \
-	     "Imported object %s from module %s is not a class" % \
-	     (name, module)
         return klass
 
+    def load_reduce(self):
+        stack = self.stack
+
+        callable = stack[-2]
+        arg_tup  = stack[-1]
+	del stack[-2:]
+
+	if (type(callable) is not ClassType):
+	    if (not safe_constructors.has_key(callable)):
+		try:
+                    safe = callable.__safe_for_unpickling__
+                except AttributeError:
+                    safe = None
+
+                if (not safe):
+                   raise UnpicklingError, "%s is not safe for " \
+                                          "unpickling" % callable
+
+        value = apply(callable, arg_tup)
+        self.append(value)
+    dispatch[REDUCE] = load_reduce
+
     def load_pop(self):
         del self.stack[-1]
     dispatch[POP] = load_pop
@@ -675,17 +776,29 @@ class Unpickler:
     dispatch[GET] = load_get
 
     def load_binget(self):
-        self.append(self.memo[`ord(self.read(1))`])
+        i = mloads('i' + self.read(1) + '\000\000\000')
+        self.append(self.memo[`i`])
     dispatch[BINGET] = load_binget
 
+    def load_long_binget(self):
+        i = mloads('i' + self.read(4))
+        self.append(self.memo[`i`])
+    dispatch[LONG_BINGET] = load_long_binget
+
     def load_put(self):
         self.memo[self.readline()[:-1]] = self.stack[-1]
     dispatch[PUT] = load_put
 
     def load_binput(self):
-        self.memo[`ord(self.read(1))`] = self.stack[-1]
+        i = mloads('i' + self.read(1) + '\000\000\000')
+        self.memo[`i`] = self.stack[-1]
     dispatch[BINPUT] = load_binput
 
+    def load_long_binput(self):
+        i = mloads('i' + self.read(4))
+        self.memo[`i`] = self.stack[-1]
+    dispatch[LONG_BINPUT] = load_long_binput
+
     def load_append(self):
         stack = self.stack
         value = stack[-1]
@@ -694,6 +807,16 @@ class Unpickler:
         list.append(value)
     dispatch[APPEND] = load_append
 
+    def load_appends(self):
+        stack = self.stack
+        mark = self.marker()
+        list = stack[mark - 1]
+	for i in range(mark + 1, len(stack)):
+            list.append(stack[i])
+
+        del stack[mark:]
+    dispatch[APPENDS] = load_appends
+           
     def load_setitem(self):
         stack = self.stack
         value = stack[-1]
@@ -703,6 +826,16 @@ class Unpickler:
         dict[key] = value
     dispatch[SETITEM] = load_setitem
 
+    def load_setitems(self):
+        stack = self.stack
+        mark = self.marker()
+        dict = stack[mark - 1]
+	for i in range(mark + 1, len(stack), 2):
+            dict[stack[i]] = stack[i + 1]
+
+        del stack[mark:]
+    dispatch[SETITEMS] = load_setitems
+
     def load_build(self):
         stack = self.stack
         value = stack[-1]